From ae970e0b25d864f57a788b576892e5e108fcefcf Mon Sep 17 00:00:00 2001 From: "Zhang, Chi" Date: Thu, 15 Oct 2026 21:14:52 -0700 Subject: [PATCH 01/10] feat(rai): move to ryzenai-corelib 0.26, add Qwen3.6-35B-A3B, run Gemma 4 at 16k Moves the `--rai` backend from ryzenai-corelib 0.11 to 0.26 and builds on what the new corelib offers. corelib 0.26 (pinned exactly): - Follow the ABI changes: `create_tensor_window` takes strides, `linear_attention_bf16` takes rows, and models are built with `model_weights_create_path` + `model_create`. - A 3-D matmul output that writes the KV cache is a strided window `[max_seq*head, head, 1]` (Phi-4, Qwen3, Qwen3.5 V writes; Gemma 4 E4B staging window). - Per-op engines open their stream through `corelib::OpenModelStream`, which enables the model's kernel bundle before any kernel query. - corelib's per-model libraries are copied beside `flm.exe` on Windows and beside corelib by `home_install.sh` on Linux. Nothing needs `RYZENAI_CORELIB_DD_ROOT` or `RYZENAI_CORELIB_KERNELS_ROOT` any more. Qwen3.6-35B-A3B (new `qwen3.6-moe-rai:35b-a3b`, text only): - Runs on corelib's model API; context up to 262144 in 4096-row pages, a prompt is a single `model_forward`, and checkpoint/restore use corelib state snapshots. - Fetch logits only for the row that gets sampled. Gemma 4: - Context of 1..4096 on flat attention, or a multiple of 4096 up to 16384 on corelib paged attention; the aie_next catalog entries default to 16384. A paged continuation starts from its own page. - `gemma4-it-rai:e4b` is offered again. Shared: - `Sampler::sampler_topk_apply` selects top-k with a k-entry heap in one pass (same ranking and tie rule). Tests: new `src/test/qwen36_rai`, and the Gemma 4, Phi-4 and Qwen3 rai tests and fake corelib follow the 0.26 ABI and the 16k paged context. Co-authored-by: Cursor --- .gitignore | 1 + docs/docs/models/phi.md | 2 +- docs/docs/models/qwen.md | 10 +- docs/rai-setup-linux.md | 59 ++-- src/CMakeLists.txt | 87 +++++- src/common/AutoModel/builtin_backends.cpp | 4 + src/common/AutoModel/modeling_qwen3_6_moe.cpp | 51 +++- src/common/models/README.md | 4 +- .../models/gemma4/rai/aie_next/gemma4_rai.cpp | 235 ++++++++++---- .../rai/aie_next/gemma4_rai_backend.cpp | 12 +- .../gemma4/rai/aie_next/gemma4_rai_gguf.cpp | 6 + .../rai/aie_next/gemma4_rai_shape_plan.cpp | 58 +++- .../models/phi4/rai/aie_next/phi4_rai.cpp | 10 +- .../models/qwen3/rai/aie_next/qwen3_rai.cpp | 12 +- .../models/qwen35/rai/aie_next/qwen35_rai.cpp | 22 +- .../models/qwen36/rai/aie_next/qwen36_rai.cpp | 289 ++++++++++++++++++ .../rai/aie_next/qwen36_rai_backend.cpp | 109 +++++++ src/common/modules/sampler.cpp | 51 ++-- src/common/rai/corelib_api.cpp | 50 ++- src/common/rai/host_ops.cpp | 17 +- src/home_install.sh | 53 ++-- .../models/gemma4/rai/aie_next/gemma4_rai.hpp | 12 +- .../rai/aie_next/gemma4_rai_backend.hpp | 20 +- .../rai/aie_next/gemma4_rai_constants.hpp | 23 +- .../gemma4/rai/aie_next/gemma4_rai_gguf.hpp | 3 + .../rai/aie_next/gemma4_rai_shape_plan.hpp | 15 +- .../phi4/rai/aie_next/phi4_rai_constants.hpp | 2 +- .../models/qwen36/rai/aie_next/qwen36_rai.hpp | 46 +++ .../rai/aie_next/qwen36_rai_backend.hpp | 25 ++ src/include/rai/corelib_api.hpp | 25 +- src/include/rai/corelib_object.hpp | 22 ++ src/include/rai/host_ops.hpp | 3 + src/include/rai/kernel_grid.hpp | 45 ++- src/model_info.json | 32 ++ src/model_list.json | 56 +++- src/test/gemma4_rai/fake_corelib.cpp | 81 ++++- src/test/gemma4_rai/fake_corelib.hpp | 17 ++ src/test/gemma4_rai/gemma4_gguf_fixture.hpp | 2 + src/test/gemma4_rai/test_gemma4_backend.cpp | 58 +++- src/test/gemma4_rai/test_gemma4_engine.cpp | 125 +++++++- src/test/gemma4_rai/test_gemma4_gguf.cpp | 18 ++ .../gemma4_rai/test_gemma4_shape_plan.cpp | 54 +++- src/test/phi4_rai/fake_corelib.cpp | 76 ++++- src/test/phi4_rai/fake_corelib.hpp | 14 + src/test/phi4_rai/run_real_rai_acceptance.ps1 | 4 +- src/test/phi4_rai/test_phi4_engine.cpp | 26 +- src/test/qwen36_rai/CMakeLists.txt | 42 +++ src/test/qwen36_rai/test_qwen36_rai.cpp | 236 ++++++++++++++ src/test/qwen3_rai/test_qwen3_rai.cpp | 27 ++ 49 files changed, 1963 insertions(+), 288 deletions(-) create mode 100644 src/common/models/qwen36/rai/aie_next/qwen36_rai.cpp create mode 100644 src/common/models/qwen36/rai/aie_next/qwen36_rai_backend.cpp create mode 100644 src/include/models/qwen36/rai/aie_next/qwen36_rai.hpp create mode 100644 src/include/models/qwen36/rai/aie_next/qwen36_rai_backend.hpp create mode 100644 src/test/qwen36_rai/CMakeLists.txt create mode 100644 src/test/qwen36_rai/test_qwen36_rai.cpp diff --git a/.gitignore b/.gitignore index 76327edd6..25b788168 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ build build_*/ +.worktrees/ .vscode __pycache__ _site diff --git a/docs/docs/models/phi.md b/docs/docs/models/phi.md index c6fae6574..f0738c4bd 100644 --- a/docs/docs/models/phi.md +++ b/docs/docs/models/phi.md @@ -74,7 +74,7 @@ The `src/test/phi4_rai` suite is still Windows-only and is not configured on Lin A rai build (`-DFLM_ENABLE_RAI=ON`) **links corelib in**, because the NPU device the whole process shares is the one behind corelib's hardware context rather than a device FastFlowLM opens itself. Point the build at the library with `RYZENAI_CORELIB_INCLUDE_DIR` and `RYZENAI_CORELIB_LIBRARY`. `FLM_RAI_CORELIB_PATH` selects a shared library (`.dll` on Windows, `.so` elsewhere) only in the older dynamically loading configuration; in a statically linked rai build it is ignored, and `flm` says so if it is set. -The corelib ABI is still pre-1.0, so FastFlowLM requires an **exact `0.11.0`** match on major, minor and patch. The version is queried before any other entry point, so a mismatched runtime reports a version error rather than a missing symbol. Within 0.11.0, FastFlowLM needs a corelib from after the weight API became one `*_weights_create` per operator (corelib commit `ab0bfee`); an older 0.11.0 header fails the build on the missing `*_weights_create` symbols. Corelib's own dependency directory must be reachable on `PATH` (Windows) or `LD_LIBRARY_PATH` (Linux). +The corelib ABI is still pre-1.0, so FastFlowLM requires an **exact `0.26.0`** match on major, minor and patch. The version is queried before any other entry point, so a mismatched runtime reports a version error rather than a missing symbol. corelib's default build installs `ryzenai_corelib` plus one `ryzenai_corelib_model_` library per model, each carrying that model's kernels; corelib loads them only from its own directory, so they must sit beside it (a Windows rai build copies them next to `flm.exe`). Each engine enables its model's kernel bundle on the stream it opens; no kernel store or environment variable is involved. Corelib's own dependency directory must be reachable on `PATH` (Windows) or `LD_LIBRARY_PATH` (Linux). ```powershell flm pull phi4-mini-it-rai:4b diff --git a/docs/docs/models/qwen.md b/docs/docs/models/qwen.md index 4cd805f2c..df513a24e 100644 --- a/docs/docs/models/qwen.md +++ b/docs/docs/models/qwen.md @@ -107,9 +107,15 @@ flm run qwen3:8b - **Backend:** `rai` — one engine for all four sizes. The GGUF's own metadata selects the size and is cross-checked against it. - **Source format:** GGUF `Q8_0`, read directly. Every weight is requantized to **group 64** while it is packed, except the 0.6B lm_head, which packs at **group 32**. - **Usable generation window:** 4095 tokens, prompt plus output. Over-capacity requests are rejected with HTTP 400 before any work reaches the device. -- **Availability:** developer build against corelib 0.11.0 — see [Phi-4 on the rai backend](phi.md). +- **Availability:** developer build against corelib 0.26.0 — see [Phi-4 on the rai backend](phi.md). -The same corelib path also runs `gemma4-it-rai:e2b`, `gemma4-it-rai:e4b` (text only) and `qwen3.5-rai:9b`. +The same corelib path also runs `gemma4-it-rai:e2b` and `gemma4-it-rai:e4b` (text only), `qwen3.5-rai:9b` and `qwen3.6-moe-rai:35b-a3b` (text only). Qwen3 and Qwen3.5 stay at 4096; the other two go further: + +- **`qwen3.6-moe-rai:35b-a3b` — up to 262144 tokens.** The default is 4096; raise it with `--ctx-len` (`-c`). Long contexts have three costs: + - **Memory, up front.** The KV cache is allocated when the model is created, at the full capacity you ask for (rounded up to whole 4096-token pages), whether or not a conversation ever fills it: about 20 KiB per token (10 attention layers × 2 KV heads × 256 × BF16, for K and V), so 262144 tokens is about 5 GiB on top of the weights. + - **Long prompts run at decode speed.** A prompt is one forward call, but only its first 4096 tokens from position 0 run as a single pass; every token after them, and every later turn's prompt, is stepped one row at a time, at about decode speed. corelib measures decode at about 26 tokens/s at a shallow context, slowing by about 2.4 ms per token for each further 4096-token page (about 18 tokens/s at 32k). So a 100k-token prompt takes on the order of an hour: about 64 minutes even at the shallow rate, and nearer two hours once the per-page slowdown is counted (an extrapolation; corelib's measurements stop at 32k). FastFlowLM also still sends a prompt in chunks of at most 4096 tokens (`max_prefill_len`), so a long prompt is a one-pass first chunk followed by stepped ones. + - **Checkpoints take host memory.** A checkpoint (what lets a conversation rewind without replaying) copies the KV cache into host memory, in whole 4096-token pages up to the current position: about 80 MiB per filled page, so up to about 5 GiB at a full 262144-token context. +- **`gemma4-it-rai:e2b` / `gemma4-it-rai:e4b` — up to 16384 tokens**, the length corelib has validated paged Gemma 4 to, and never past the GGUF's own `gemma4.context_length`. The default is 16384. Up to 4096 any length works; above 4096 the context must be a multiple of 4096. The KV cache for the whole context is allocated up front, when the model loads. ```powershell flm pull qwen3-rai:4b diff --git a/docs/rai-setup-linux.md b/docs/rai-setup-linux.md index 7f77d7a07..4a7f7ad8a 100644 --- a/docs/rai-setup-linux.md +++ b/docs/rai-setup-linux.md @@ -8,10 +8,10 @@ repository, so most of the work is standing up that prefix. Written against the working install on this machine. Paths use `$P` for the dependency prefix (here `/home/alfxu/ddprefix`); substitute your own. -> **Status:** at corelib 0.11.0 the rai backend has been run on hardware on -> Windows only. The Linux build configures and links against this layout, but -> it has not been run end to end on Linux at this revision. The -> [Known issues](#known-issues) are from an earlier Linux bring-up. +> **Status:** the rai backend has been run on hardware on Windows only. The +> steps below are updated for corelib 0.26.0, but they have not been run end +> to end on Linux at this revision. The [Known issues](#known-issues) are from +> an earlier Linux bring-up. --- @@ -29,48 +29,44 @@ step 4 and installing XRT itself. --- -## 1. Dependency prefix: DynamicDispatch, then corelib +## 1. Dependency prefix: corelib -This is the hard part and it is already written up in **`$P/BUILD-NOTES.md`** — -read that first. It records the two things that actually matter (the required -DD branch, and `BUILD_SHARED_LIBS=ON` for DD), the sudo-free dependency prefix, -the non-PIC link traps, and the exact `cmake` command lines for both projects. - -Do not re-derive it. The short version: +corelib no longer links DynamicDispatch, so there is no DD prefix to +build first. **`$P/BUILD-NOTES.md`** still records the sudo-free dependency +prefix and the non-PIC link traps; its DynamicDispatch steps no longer apply. ```bash P=/home/alfxu/ddprefix # your dependency prefix -# DynamicDispatch -> $P/install-dds (see BUILD-NOTES.md for the full line) -cmake -S DynamicDispatch -B build-dds ... -DCMAKE_INSTALL_PREFIX=$P/install-dds -cmake --build build-dds -j"$(nproc)" --target install - # ryzenai-corelib -> $P/install-corelib cmake -S ryzenai-corelib -B build-corelib -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_INSTALL_PREFIX=$P/install-corelib \ - -DRYZENAI_CORELIB_STAGE_RUNTIME_DLLS=OFF \ - -DCMAKE_PREFIX_PATH="$P/install-dds;$P/shim;$P/root/usr;$P/root/usr/lib/x86_64-linux-gnu;/opt/xilinx/xrt" -cmake --build build-corelib -j"$(nproc)" + -DCMAKE_PREFIX_PATH="$P/shim;$P/root/usr;$P/root/usr/lib/x86_64-linux-gnu;/opt/xilinx/xrt" +cmake --build build-corelib -j"$(nproc)" --target install ``` +The default build installs `libryzenai_corelib.so` and one +`libryzenai_corelib_model_.so` per model, each carrying that model's +kernels. corelib loads a model library only from the directory it was itself +loaded from, so those libraries must sit beside it; nothing else is needed at +run time, no kernel store and no environment variable. + **Version must match exactly.** `src/include/rai/corelib_api.hpp` has a -`#error` for anything other than corelib **0.11.0**: +`#error` for anything other than corelib **0.26.0**: ```bash grep -E 'CORELIB_VERSION_(MAJOR|MINOR|PATCH)' \ $P/install-corelib/include/ryzenai/corelib.h ``` -Must print `0`, `11`, `0`. (`BUILD-NOTES.md` still says 0.4.0 in its opening -line — that text is stale; the installed headers are 0.11.0.) The headers must -also be from after the weight API became one `*_weights_create` per operator; -an older 0.11.0 checkout fails the FastFlowLM build on those symbols. +Must print `0`, `26`, `0`. (`BUILD-NOTES.md` still says 0.4.0 in its opening +line — that text is stale.) ### Verify corelib before going further ```bash source /opt/xilinx/xrt/setup.sh -export LD_LIBRARY_PATH=$P/install-dds/lib:$P/install-corelib/lib:/opt/xilinx/xrt/lib:$LD_LIBRARY_PATH +export LD_LIBRARY_PATH=$P/install-corelib/lib:/opt/xilinx/xrt/lib:$LD_LIBRARY_PATH $P/build-corelib/tests/ryzenai_corelib_tests ``` @@ -98,8 +94,9 @@ What the flags do: - `--rai` selects the `linux-rai-on` preset (`-DFLM_ENABLE_RAI=ON`) and installs to `/scratch/$USER/flm_exe_rai` rather than `~/flm_exe`. A rai prefix stages - corelib and DynamicDispatch, which are far too large for a home directory, and - a rai build and a stock one would otherwise overwrite each other. + corelib and every `libryzenai_corelib_model_*.so` beside it, whose model + kernels are far too large for a home directory, and a rai build and a stock + one would otherwise overwrite each other. - `FLM_PREFIX=/path` overrides the install location. - `--no-build` installs an existing `build/` tree without rebuilding. @@ -121,10 +118,11 @@ from the corelib *build tree*, not a stripped or patched copy: ```bash objdump -p src/lib/libryzenai_corelib.so | grep RUNPATH -# -> /home/alfxu/ddprefix/install-dds/lib:/opt/xilinx/xrt/lib: +# -> should name /opt/xilinx/xrt/lib ``` -Without that RUNPATH it will not find DynamicDispatch at run time. +Without that RUNPATH it will not find XRT at run time. Copy every +`libryzenai_corelib_model_*.so` into `src/lib/` beside it too. ### Rebuilding after a source change @@ -150,7 +148,7 @@ rm -f /scratch/$USER/flm_exe_rai/lib/libxrt_coreutil.so.2 ``` This is a known bug in the installer's `ldd` staging walk -(`home_install.sh:251-261`), not yet fixed. A rai prefix deliberately bundles no +(`home_install.sh:247-256`), not yet fixed. A rai prefix deliberately bundles no XRT — `flm_env.sh` pins `XILINX_XRT` to the real root instead — so the staged copy is wrong by construction. @@ -239,6 +237,9 @@ shape is a dead end. | `RYZENAI_CORELIB_SAFE_MODE` | force staging buffers, no host-pointer import | | `RYZENAI_CORELIB_NO_RUN_POOL` | disable run-object reuse | | `RYZENAI_CORELIB_NO_CRASH_HANDLER` | leave crashes to the debugger | +| `RYZENAI_CORELIB_FALLBACK_PROFILE=1` | when a stream is released, report host (CPU fallback) time per artifact; no output means nothing fell back | +| `RYZENAI_CORELIB_MAX_RUNS_IN_FLIGHT` | how many runs of one kernel a stream lets run ahead (64 unset, `0` for no cap) | +| `RYZENAI_CORELIB_KERNELS_ROOT` | a kernel store to load from, for corelib development only; every model library already carries its kernels, so a normal run sets nothing | ### Other rough edges diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c53826f16..eb96e8809 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -357,6 +357,61 @@ if(FLM_ENABLE_RAI) # to resolve even though nothing calls them here. target_link_libraries(flm_rai PUBLIC "${RYZENAI_CORELIB_LIBRARY}" ${CMAKE_DL_LIBS}) target_link_libraries(flm PRIVATE flm_rai) + # corelib loads each ryzenai_corelib_model_.dll, which carries that + # model's kernels, only from the directory ryzenai_corelib.dll itself was + # loaded from, so both have to land beside flm.exe. + if(WIN32) + get_filename_component(FLM_CORELIB_LIB_DIR "${RYZENAI_CORELIB_LIBRARY}" DIRECTORY) + file(GLOB FLM_CORELIB_RUNTIME_DLLS + "${FLM_CORELIB_LIB_DIR}/../bin/ryzenai_corelib.dll" + "${FLM_CORELIB_LIB_DIR}/../bin/ryzenai_corelib_model_*.dll") + list(FILTER FLM_CORELIB_RUNTIME_DLLS EXCLUDE REGEX "ryzenai_corelib_model_dummy\\.dll$") + if(FLM_CORELIB_RUNTIME_DLLS) + add_custom_command(TARGET flm POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different + ${FLM_CORELIB_RUNTIME_DLLS} $) + else() + message(WARNING "No ryzenai_corelib DLLs under ${FLM_CORELIB_LIB_DIR}/../bin; " + "copy them beside flm.exe by hand") + endif() + endif() +endif() + +# ——————————————————————————————————————————————— +# Opt-in Windows dependency prefix. +# +# Windows otherwise has exactly two routes for Boost/curl/FFmpeg: a vcpkg +# toolchain, or the bare CI runner's hardcoded C:/dev paths and CI-specific +# library names. A developer box with neither can point this at one prefix +# laid out as /include + /lib -- in practice a conda env's +# Library dir (conda-forge boost, libcurl, ffmpeg) -- and each block that +# applies the CI runner's include dirs, definitions, link dirs and library +# names switches to it. +# Empty (the default) leaves both existing routes exactly as they were. +# FFTW is not taken from the prefix on any route: its header (include/fftw3.h) +# and import libs (lib/libfftw3*-3.lib) are vendored in this tree. +# ——————————————————————————————————————————————— +set(FLM_WIN_DEPS_PREFIX "" CACHE PATH + "Windows only: take Boost/curl/FFmpeg from /include and /lib (e.g. a conda env's Library dir) instead of the CI runner's C:/dev paths") +if(FLM_WIN_DEPS_PREFIX) + if(NOT WIN32) + message(FATAL_ERROR "FLM_WIN_DEPS_PREFIX is a Windows-only route") + endif() + if(VCPKG_TOOLCHAIN) + message(FATAL_ERROR "FLM_WIN_DEPS_PREFIX and a vcpkg toolchain are two " + "different dependency routes; use one") + endif() + # The library names a conda-forge Library dir ships: unversioned Boost and + # curl import libs (both are DLL builds), FFmpeg import libs. + set(FLM_WIN_DEPS_PREFIX_LIBS + avformat avcodec avutil swscale swresample libcurl boost_program_options) + foreach(_flm_dep IN LISTS FLM_WIN_DEPS_PREFIX_LIBS) + if(NOT EXISTS "${FLM_WIN_DEPS_PREFIX}/lib/${_flm_dep}.lib") + message(FATAL_ERROR "FLM_WIN_DEPS_PREFIX: missing " + "${FLM_WIN_DEPS_PREFIX}/lib/${_flm_dep}.lib") + endif() + endforeach() + message(STATUS "Windows deps from prefix: ${FLM_WIN_DEPS_PREFIX}") endif() if(WIN32) @@ -373,6 +428,8 @@ if(WIN32) find_package(FFTW3f CONFIG REQUIRED) find_package(FFTW3l CONFIG REQUIRED) endif() + # A third, opt-in route (FLM_WIN_DEPS_PREFIX, declared above) takes Boost, + # curl and FFmpeg from one local prefix, wired next to each CI-runner block. # Otherwise (the bare self-hosted CI runner) there is no vcpkg toolchain and # no Boost in vcpkg: Boost is a standalone b2 build under C:/dev/boost_1_88_0 # and curl/ffmpeg/fftw are linked by raw name from the vcpkg lib dir. That @@ -424,12 +481,14 @@ if(NOT FLM_USE_HRX) endif() endif() -if(WIN32 AND NOT VCPKG_TOOLCHAIN) +if(WIN32 AND NOT VCPKG_TOOLCHAIN AND NOT FLM_WIN_DEPS_PREFIX) # Bare CI runner: standalone Boost + vcpkg headers by absolute path. target_include_directories(flm PUBLIC C:/dev/boost_1_88_0 C:/dev/vcpkg/installed/x64-windows/include/ ) +elseif(FLM_WIN_DEPS_PREFIX) + target_include_directories(flm PUBLIC "${FLM_WIN_DEPS_PREFIX}/include") endif() target_compile_definitions(flm PUBLIC @@ -456,13 +515,22 @@ if(WIN32) ) endif() -if(WIN32 AND NOT VCPKG_TOOLCHAIN) +if(WIN32 AND NOT VCPKG_TOOLCHAIN AND NOT FLM_WIN_DEPS_PREFIX) # Bare CI runner: statically link the standalone Boost + curl. target_compile_definitions(flm PUBLIC CURL_STATICLIB BOOST_ALL_NO_LIB BOOST_ALL_STATIC_LINK ) +elseif(FLM_WIN_DEPS_PREFIX) + # conda-forge Boost and curl are DLL builds. program_options exports data + # (boost::program_options::arg, options_description::m_default_line_length) + # that only resolves through dllimport, so BOOST_ALL_STATIC_LINK would leave + # it unresolved at link time; CURL_STATICLIB would mis-declare curl likewise. + target_compile_definitions(flm PUBLIC + BOOST_ALL_NO_LIB + BOOST_ALL_DYN_LINK + ) endif() target_link_directories(flm PUBLIC @@ -480,12 +548,14 @@ if(NOT FLM_USE_HRX) endif() endif() -if(WIN32 AND NOT VCPKG_TOOLCHAIN) +if(WIN32 AND NOT VCPKG_TOOLCHAIN AND NOT FLM_WIN_DEPS_PREFIX) # Bare CI runner: standalone Boost stage libs + vcpkg import libs. target_link_directories(flm PUBLIC C:/dev/boost_1_88_0/stage/lib C:/dev/vcpkg/installed/x64-windows/lib ) +elseif(FLM_WIN_DEPS_PREFIX) + target_link_directories(flm PUBLIC "${FLM_WIN_DEPS_PREFIX}/lib") endif() # Link static libraries first @@ -602,6 +672,17 @@ if(WIN32 AND VCPKG_TOOLCHAIN) FFTW3::fftw3f FFTW3::fftw3l ) +elseif(WIN32 AND FLM_WIN_DEPS_PREFIX) + # Local prefix: the names it ships (checked at configure time above). FFTW + # is the tree's vendored lib/libfftw3*-3, as on the CI runner. flm itself + # calls only fftwf_* and no prebuilt engine DLL imports FFTW, so the double + # and long-double libs contribute no import; they stay for parity. + target_link_libraries(flm PUBLIC + ${FLM_WIN_DEPS_PREFIX_LIBS} + libfftw3-3 + libfftw3f-3 + libfftw3l-3 + ) elseif(WIN32) # Bare CI runner: link the native deps by raw name from the standalone Boost # + vcpkg lib dirs wired above (no CMake package configs there). diff --git a/src/common/AutoModel/builtin_backends.cpp b/src/common/AutoModel/builtin_backends.cpp index 6f2b77089..3bc66a219 100644 --- a/src/common/AutoModel/builtin_backends.cpp +++ b/src/common/AutoModel/builtin_backends.cpp @@ -13,6 +13,7 @@ #include "models/phi4/rai/aie_next/phi4_rai_backend.hpp" #include "models/qwen3/rai/aie_next/qwen3_rai_backend.hpp" #include "models/qwen35/rai/aie_next/qwen35_rai_backend.hpp" +#include "models/qwen36/rai/aie_next/qwen36_rai_backend.hpp" #endif namespace flm::backend { @@ -70,6 +71,9 @@ void register_builtin_backends(BackendRegistry& registry) { registry.register_backend("qwen3.5", flm::backend::kRaiBackendId, flm::qwen35::rai_factory(), flm::qwen35::rai_traits()); + registry.register_backend("qwen3.6-moe", flm::backend::kRaiBackendId, + flm::qwen36::rai_factory(), + flm::qwen36::rai_traits()); #endif } diff --git a/src/common/AutoModel/modeling_qwen3_6_moe.cpp b/src/common/AutoModel/modeling_qwen3_6_moe.cpp index 4c00fe310..0a3a98d8c 100644 --- a/src/common/AutoModel/modeling_qwen3_6_moe.cpp +++ b/src/common/AutoModel/modeling_qwen3_6_moe.cpp @@ -59,6 +59,15 @@ bool Qwen3_6_MOE::insert(chat_meta_info_t& meta_info, lm_uniform_input_t& input, header_print("WARNING", "No messages or prompt provided"); return false; } + if (this->backend_id() == flm::backend::kRaiBackendId) { + bool has_image = !input.images.empty(); + for (const auto& item : input.messages) { + if (item.contains("images") && !item["images"].empty()) has_image = true; + } + if (has_image) + throw ModelRequestError(400, false, + "Qwen3.6 on backend 'rai' runs text only; this request has images"); + } constexpr bool DEBUG_IMAGE_PREPROCESS = false; qwen3_6_moe_image_payload_t image_payload; @@ -279,11 +288,10 @@ bool Qwen3_6_MOE::insert(chat_meta_info_t& meta_info, lm_uniform_input_t& input, // hardware int restore_idx = -1; - qwen3_6_moe_npu *qwen3_6_moe_engine = dynamic_cast(this->lm_engine); const bool has_images = image_payload.num_images > 0; if (meta_info.restore_allowed) { - restore_idx = qwen3_6_moe_engine->restore(); + restore_idx = this->lm_engine->restore(); this->total_tokens = restore_idx; this->token_history = checkpoint_his; // restore the token history to be consistent with the restored KV cache, which is crucial for correct functioning of _shared_insert's prefix-matching logic } @@ -296,7 +304,7 @@ bool Qwen3_6_MOE::insert(chat_meta_info_t& meta_info, lm_uniform_input_t& input, : this->_shared_insert(meta_info, tokens, is_cancelled, nullptr); checkpoint_his = token_history; - int checkpoint_idx = qwen3_6_moe_engine->checkpoint(); + int checkpoint_idx = this->lm_engine->checkpoint(); return success; } @@ -317,6 +325,16 @@ std::string Qwen3_6_MOE::generate(chat_meta_info_t& meta_info, int length_limit, this->profiler_list[TKOEN_DECODE_TIME].reset(); std::string token_str; int sampled_token; + // Some backends refuse to decode past a limit of their own, below MAX_L. + const uint32_t decode_cap = this->decode_cap(); + // insert() trimmed the opening of the think block; it is decoded here, and + // every one of those steps occupies the context. + const uint32_t forced = this->enable_think ? 2 : 4; + if (this->total_tokens + forced >= decode_cap) { + header_print("WARNING", "Max length reached, stopping generation..."); + meta_info.stop_reason = MAX_LENGTH_REACHED; + return result; + } if(this->enable_think) { this->token_history.push_back(think_start_id); this->profiler_list[DECODING_TIME].start(); @@ -366,19 +384,20 @@ std::string Qwen3_6_MOE::generate(chat_meta_info_t& meta_info, int length_limit, this->_apply_tool_choice_mask(y, meta_info); sampled_token = this->sampler->sample(y); } - this->total_tokens++; + this->total_tokens += forced; meta_info.generated_tokens++; int last_sampled_token = sampled_token; token_str = this->tokenizer->run_time_decoder(last_sampled_token); result += token_str; os << token_str << std::flush; - if (this->total_tokens >= this->MAX_L){ - header_print("WARNING", "Max length reached, stopping generation..."); - reason = MAX_LENGTH_REACHED; - return result; - } - while (this->total_tokens < this->MAX_L){ + while (true){ + // Running out of window is a truncation, not the model choosing to + // stop; testing here rather than in the loop condition records it. + if (this->total_tokens >= decode_cap){ + reason = MAX_LENGTH_REACHED; + break; + } if (is_cancelled()) { reason = CANCEL_DETECTED; // reset stream content @@ -411,7 +430,10 @@ std::string Qwen3_6_MOE::generate(chat_meta_info_t& meta_info, int length_limit, this->token_history.push_back(sampled_token); if (this->is_eos(sampled_token)){ meta_info.generated_tokens++; - this->lm_engine->forward(last_sampled_token); + if (this->forward_on_eos && + (!this->backend_ || this->backend_->forwards_past_eos())) { + this->lm_engine->forward(last_sampled_token); + } break; } meta_info.generated_tokens++; @@ -422,7 +444,7 @@ std::string Qwen3_6_MOE::generate(chat_meta_info_t& meta_info, int length_limit, } meta_info.decoding_duration = (uint64_t)(time_utils::cast_to_us(this->profiler_list[DECODING_TIME].get_total_time()).first) * 1e3; meta_info.stop_reason = reason; - if (this->total_tokens >= this->MAX_L){ + if (this->total_tokens >= decode_cap){ header_print("WARNING", "Max length reached, stopping generation..."); } @@ -437,9 +459,8 @@ std::string Qwen3_6_MOE::generate_with_prompt(chat_meta_info_t& meta_info, lm_un return ""; } header_print("FLM", "Prompt inserted, starting generation..."); - qwen3_6_moe_npu* qwen36_engine = dynamic_cast(this->lm_engine); - int checkpoint_idx = qwen36_engine->checkpoint(); - int restore_idx = qwen36_engine->restore(); + int checkpoint_idx = this->lm_engine->checkpoint(); + int restore_idx = this->lm_engine->restore(); header_print_r("FLM", "Checkpoint before generation: " << checkpoint_idx << ", restore point: " << restore_idx << ", user context length: " << this->token_history.size()); if (this->enable_think) { os << "\n" << std::flush; diff --git a/src/common/models/README.md b/src/common/models/README.md index bfb2545bd..9d5d1922c 100644 --- a/src/common/models/README.md +++ b/src/common/models/README.md @@ -35,7 +35,7 @@ This is a contributor document. For *using* a backend once it exists — `--back |---|---| | **Platform** | Windows and Linux. Both configure and build; only Windows has been run on hardware, and the `src/test/phi4_rai` suite is still Windows-only. | | **Hardware** | An aie_next NPU. There is no simulator; a wrong shape shows up as garbage output, not an error. | -| **corelib headers** | Exactly **0.11.0**, from after the weight API became one `*_weights_create` per operator. [`corelib_api.hpp`](../../include/rai/corelib_api.hpp) `#error`s on any other version — deliberately, because the C ABI has changed shape between patch releases. | +| **corelib headers** | The corelib ABI is still pre-1.0, so FastFlowLM requires an **exact `0.26.0`** match on major, minor and patch. [`corelib_api.hpp`](../../include/rai/corelib_api.hpp) `#error`s on any other version — deliberately, because the C ABI has changed shape between patch releases. The version is queried before any other entry point, so a mismatched runtime reports a version error rather than a missing symbol. corelib's default build gives one `ryzenai_corelib_model_` library per model, carrying its kernels, which must sit beside corelib. Each engine enables its model's kernel bundle on the stream it opens. | | **Weights** | A GGUF the vendor kernels can requantize. Phi-4 uses Q8_0; each weight is a `*_weights_create` call whose qweight source states `gguf_q8_0`. | Configure with: @@ -166,7 +166,7 @@ engine type, never through a `causal_lm*`. mapped range and produces its own object — so they run across a pool (`kWeightCreateConcurrency`). The per-create thread hint (`kRequantizeThreads`) goes to corelib as `fast_packer_threads`: 0 lets corelib - choose (8 threads at 0.11.0), 1 keeps each create serial. The two multiply, + choose (8 threads), 1 keeps each create serial. The two multiply, so pick them together. See [`phi4_rai_constants.hpp`](../../include/models/phi4/rai/aie_next/phi4_rai_constants.hpp). - Expose `bool poisoned() const noexcept`. A corelib failure mid-decode usually diff --git a/src/common/models/gemma4/rai/aie_next/gemma4_rai.cpp b/src/common/models/gemma4/rai/aie_next/gemma4_rai.cpp index 49fba178d..8b1671f21 100644 --- a/src/common/models/gemma4/rai/aie_next/gemma4_rai.cpp +++ b/src/common/models/gemma4/rai/aie_next/gemma4_rai.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -194,6 +195,13 @@ struct gemma4_rai::Impl { std::shared_ptr api; Gemma4Config shape; std::uint32_t max_length; + /// \brief context longer than one flat_mha window: paged KV, PDI (8, 33) + bool paged{}; + /// \brief rows the KV cache and rotary tables hold + /// \note Unpaged this stays kMaxSequenceLength even when the caller asked + /// for less, which is the pitch flat_mha addresses. Paged it is the + /// constructed context, a whole number of pages. + std::uint32_t cache_rows{}; int position{}; std::optional saved; bool poisoned{}; @@ -328,9 +336,19 @@ struct gemma4_rai::Impl { api(runtime ? runtime->api() : nullptr), max_length(maximum) { if (!package) throw std::invalid_argument("Gemma 4 GGUF package is null"); if (!runtime || !api) throw std::invalid_argument("corelib runtime is null"); - if (!maximum || maximum > kMaxSequenceLength) - throw std::invalid_argument("Gemma 4 maximum length must be in 1..4096"); + if (!ContextFits(maximum)) { + throw std::invalid_argument( + "Gemma 4 maximum length must be in 1..4096, or a multiple of 4096 up to " + + std::to_string(kMaxPagedContext)); + } + paged = maximum > static_cast(kPageSize); + cache_rows = paged ? maximum : static_cast(kMaxSequenceLength); shape = shape_override ? std::move(*shape_override) : package->Config(); + if (paged && static_cast(maximum) > shape.context_length) { + throw std::invalid_argument( + "Gemma 4 maximum length " + std::to_string(maximum) + + " is past this GGUF's context of " + std::to_string(shape.context_length)); + } const auto& cfg = shape; const auto& file = package->File(); @@ -526,11 +544,11 @@ struct gemma4_rai::Impl { // the ELF set, and a shape exists under one pair and not another), so // the plan cannot exist before a stream does. auto lease = runtime->AcquireExecution(); - void* raw_stream = nullptr; - api->Check(api->functions().create_stream(kPrefillPdi, kTokenPdi, &raw_stream), - "ryzenai_corelib_create_stream"); - stream = UniqueStream(api, raw_stream); - plan = Gemma4ShapePlan::Build(api, stream.get(), cfg); + const int token_pdi = paged ? kPagedTokenPdi : kTokenPdi; + stream = corelib::OpenModelStream(api, kPrefillPdi, token_pdi, "gemma4"); + plan = Gemma4ShapePlan::Build(api, stream.get(), cfg, + paged ? static_cast(maximum) + : kMaxSequenceLength); // ---- the slot table std::vector slots; @@ -804,8 +822,14 @@ struct gemma4_rai::Impl { // THE WIDTH IS THE OWNING LAYER'S OWN HEAD. A layer below // `kv_layers` owns cache `index`, so this is layer `cache`'s head. const auto head = cfg.LayerHeadDim(cache); - k_cache.push_back(bf16({cfg.kv_heads, rows, head}, "K cache")); - v_cache.push_back(bf16({cfg.kv_heads, rows, head}, "V cache")); + if (paged) { + const auto pages = static_cast(cache_rows) / kPageSize; + k_cache.push_back(bf16({pages, cfg.kv_heads, kPageSize, head}, "K cache")); + v_cache.push_back(bf16({pages, cfg.kv_heads, kPageSize, head}, "V cache")); + } else { + k_cache.push_back(bf16({cfg.kv_heads, rows, head}, "K cache")); + v_cache.push_back(bf16({cfg.kv_heads, rows, head}, "V cache")); + } } } @@ -837,7 +861,8 @@ struct gemma4_rai::Impl { const auto& cfg = shape; const auto view = [&](UniqueHostView& destination, std::int64_t head_dim, const std::vector& values, const char* label) { - const std::array dims{kMaxSequenceLength, head_dim / 2}; + const std::array dims{static_cast(cache_rows), + head_dim / 2}; void* created = nullptr; api->Check(api->functions().create_host_view( ryzenai_corelib_data_type_fp32, dims.data(), dims.size(), @@ -846,14 +871,15 @@ struct gemma4_rai::Impl { destination = UniqueHostView(api, created); }; // Base `rope_theta_swa` (1e4) over the sliding head, NO factors. - auto sliding = MakeRopeTables(cfg.head_dim, cfg.rope_theta_swa, kMaxSequenceLength); + const auto rope_rows = static_cast(cache_rows); + auto sliding = MakeRopeTables(cfg.head_dim, cfg.rope_theta_swa, rope_rows); rope_cos_sliding = std::move(sliding.cos); rope_sin_sliding = std::move(sliding.sin); view(cos_sliding, cfg.head_dim, rope_cos_sliding, "sliding cos"); view(sin_sliding, cfg.head_dim, rope_sin_sliding, "sliding sin"); // Base `rope_theta` (1e6) over the full head, WITH `rope_freqs.weight` // as a per-pair frequency DIVISOR. - auto full = MakeRopeTables(cfg.global_head_dim, cfg.rope_theta, kMaxSequenceLength, + auto full = MakeRopeTables(cfg.global_head_dim, cfg.rope_theta, rope_rows, factors.values); rope_cos_full = std::move(full.cos); rope_sin_full = std::move(full.sin); @@ -862,10 +888,10 @@ struct gemma4_rai::Impl { } /// \brief run every slot's `pack()` across a small pool - /// \note TWO FORMS OF PARALLELISM THAT MUST NOT MULTIPLY. Each create is - /// submitted at `kRequantizeThreads` (0 -- corelib's default of - /// one) while `kWeightCreateConcurrency` of them run at once, so - /// the total is the pool size and not the product. The reference + /// \note TWO FORMS OF PARALLELISM THAT MULTIPLY. Each create is + /// submitted at `kRequantizeThreads` (0 -- corelib picks, 8 packer + /// threads at 0.25.0) while `kWeightCreateConcurrency` of them run + /// at once, so up to the product runs together. The reference /// driver packs strictly serially and says why (an open, /// unattributed corelib defect correlating concurrent creates with /// all-zero device output); this follows Phi-4's shape instead, @@ -1001,10 +1027,12 @@ struct gemma4_rai::Impl { /// \brief create a window onto `parent`, naming it in any failure UniqueTensorWindow Window(void* parent, std::vector dims, - std::size_t offset, const char* label) { + std::size_t offset, const char* label, + std::vector strides = {}) { void* created = nullptr; - api->Check(api->functions().create_tensor_window(parent, dims.data(), dims.size(), - offset, &created), + api->Check(api->functions().create_tensor_window(parent, dims.data(), dims.size(), offset, + strides.empty() ? nullptr : strides.data(), + strides.size(), &created), std::string("ryzenai_corelib_create_tensor_window ") + label); return UniqueTensorWindow(api, created); } @@ -1021,8 +1049,10 @@ struct gemma4_rai::Impl { /// a supported shape costs. /// /// DECODE IS NOT PADDED: an M == 1 kernel ships for every operator here. - static std::int64_t PaddedRowsFor(std::int64_t rows) { + std::int64_t PaddedRowsFor(std::int64_t rows) const { if (rows <= 1) return 1; + // A paged chunk past page 0 is only correct at a full page. + if (paged && position > 0 && rows <= kPageSize) return kPageSize; for (const std::int64_t bucket : {128, 256, 512, 1024, 2048, 3072, 4096}) { if (rows <= bucket) return bucket; } @@ -1030,6 +1060,18 @@ struct gemma4_rai::Impl { " rows: the largest shipped prefill kernel is 4096"); } + /// \brief byte-element offset of one KV row inside its cache + std::size_t CacheOffset(std::int64_t kv_head, std::int64_t index, + std::int64_t head) const { + if (!paged) { + return static_cast((kv_head * kMaxSequenceLength + index) * head); + } + const auto page = index / kPageSize; + const auto row = index % kPageSize; + return static_cast( + ((page * shape.kv_heads + kv_head) * kPageSize + row) * head); + } + /// \brief every operator agrees that `padded` is the extent it runs at /// /// ONE ROW EXTENT FOR THE WHOLE PASS, AND IT HAS TO BE CHECKED RATHER THAN @@ -1113,11 +1155,12 @@ struct gemma4_rai::Impl { // WHERE V'S PROJECTION LANDS, AND IT IS HEAD-MAJOR BECAUSE THE CACHE // IS. A cache is [kv_heads, max_seq, head]; a projection's ordinary // output is [rows, kv_heads * head], token-major. Those are the same - // bytes only when there is ONE KV head. With more, the norm would have - // to read a head's rows at a stride, and a window is an offset and a - // shape with no stride at all -- so the staging buffer has to be - // written head-major in the first place, which only the scattering - // kernel can do, and a 3-D output is what asks for it. + // bytes only when there is ONE KV head. With more, one head's rows sit + // `kv_heads * head` apart in the token-major output, while V's norm + // below reads each head as a dense [rows, head] plane at an offset -- + // so the staging buffer has to be written head-major in the first + // place. Only the scattering kernel can do that, and a 3-D output, a + // window strided [max_seq * head, head, 1], is what asks for it. // // AT EXACTLY ONE KV HEAD A PLAIN 2-D OUTPUT IS USED INSTEAD, and that // is not merely an optimization: with one head there is nothing to @@ -1128,7 +1171,8 @@ struct gemma4_rai::Impl { geometry.v_dest = cfg.kv_heads == 1 ? Window(v_buf.get(), {padded, head}, 0, "v") - : Window(v_buf.get(), {cfg.kv_heads, kMaxSequenceLength, head}, 0, "v"); + : Window(v_buf.get(), {cfg.kv_heads, kMaxSequenceLength, head}, 0, "v", + {kMaxSequenceLength * head, head, 1}); // QK-Norm normalizes ONE HEAD, so its row count is tokens x heads. geometry.q_heads_in = Window(q_buf.get(), {padded * cfg.q_heads, head}, 0, "q heads"); @@ -1258,21 +1302,23 @@ struct gemma4_rai::Impl { // the reference driver has this one and not the other, and a future // caller that decouples them would otherwise lose the only check that // is about the thing the kernel actually addresses. - if (position < 0 || position + padded > kMaxSequenceLength) { + const auto cache_limit = static_cast(cache_rows); + if (position < 0 || position + padded > cache_limit) { throw std::out_of_range( std::to_string(rows) + " rows at position " + std::to_string(position) + " run a " + std::to_string(padded) + "-row kernel, which reaches row " + std::to_string(position + padded) + " of a " + - std::to_string(kMaxSequenceLength) + "-row cache"); + std::to_string(cache_limit) + "-row cache"); } - // THERE IS NO CHUNKED PREFILL. `flat_mha_bf16` rejects a multi-row - // call at a non-zero position itself -- but as the SIXTH dispatch of a - // layer, which is the abort described above rather than an exception. - if (rows > 1 && position != 0) { + // Chunked prefill is paged only, and only on a page boundary. + // flat_mha rejects a multi-row call at a non-zero position itself, but + // as the sixth dispatch of a layer, which aborts the process. + const bool page_chunk = paged && position > 0 && position % kPageSize == 0; + if (rows > 1 && position != 0 && !page_chunk) { throw std::invalid_argument( std::to_string(rows) + " rows at position " + std::to_string(position) + - ": attention ships no chunked prefill, so a multi-row pass must start " - "at position 0 and a continuation is one row at a time"); + ": a multi-row pass must start at position 0" + + (paged ? " or on a 4096-row page boundary" : "")); } RequireOneRowExtent(padded); @@ -1345,9 +1391,7 @@ struct gemma4_rai::Impl { for (std::int64_t kv_head = 0; kv_head < cfg.kv_heads; ++kv_head) { windows.push_back(Window( v_cache[static_cast(layer)].get(), {padded, head}, - static_cast(kv_head * kMaxSequenceLength * head + - position * head), - "V cache")); + CacheOffset(kv_head, position, head), "V cache")); } } @@ -1432,20 +1476,27 @@ struct gemma4_rai::Impl { // TWO ROTARY PAIRS, AND `swa` SELECTS. One pair used // everywhere still converges to fluent text. const auto owner = static_cast(cfg.layer_cache_owner[slot]); - dispatch(api->functions().flat_mha( - stream.get(), - // SELECT, NEVER MUTATE. The plan holds four - // immutable descriptors keyed (geometry, cache - // role); there is deliberately no one-argument - // accessor, so a loop that forgot the cache role - // would not compile. - &plan.attention_desc(swa, /*shared=*/!owns), - kind.q_attn.get(), kind.k_attn.get(), position, - (swa ? cos_sliding : cos_full).get(), - (swa ? sin_sliding : sin_full).get(), - k_cache[owner].get(), v_cache[owner].get(), - kind.attn.get()), - "ryzenai_corelib_flat_mha_bf16", index); + void* cos = (swa ? cos_sliding : cos_full).get(); + void* sin = (swa ? sin_sliding : sin_full).get(); + if (plan.paged()) { + dispatch(api->functions().paged_mha( + stream.get(), &plan.paged_attention_desc(swa, /*shared=*/!owns), + kind.q_attn.get(), kind.k_attn.get(), position, cos, sin, + k_cache[owner].get(), v_cache[owner].get(), kind.attn.get()), + "ryzenai_corelib_paged_mha_bf16", index); + } else { + dispatch(api->functions().flat_mha( + stream.get(), + // SELECT, NEVER MUTATE. The plan holds four + // immutable descriptors keyed (geometry, cache + // role); there is deliberately no one-argument + // accessor, so a loop that forgot the cache role + // would not compile. + &plan.attention_desc(swa, /*shared=*/!owns), + kind.q_attn.get(), kind.k_attn.get(), position, cos, sin, + k_cache[owner].get(), v_cache[owner].get(), kind.attn.get()), + "ryzenai_corelib_flat_mha_bf16", index); + } dispatch(api->functions().matmul(stream.get(), kind.attn.get(), weights.o.get(), hidden_w.get()), "ryzenai_corelib_matmul_bf16 output", index); @@ -1597,16 +1648,66 @@ struct gemma4_rai::Impl { // continuation that no longer exists. Every accessor is behind // `usable()`, so this is defence rather than something a caller // can observe. + // + // ONE EXCEPTION. A paged prefill chunk past page 0 can come back + // non-finite from an intermittent kernel defect: the chunk + // rewrites only its own page, and `position` has not moved, so + // RunPrompt may run that same chunk again. Unpaged passes and + // page 0 still poison. The message text is RequireFinite's. + bool retry = false; + try { + throw; + } catch (const std::exception& error) { + retry = paged && position > 0 && + std::string(error.what()).find("non-finite") != std::string::npos; + } catch (...) { + } if (submitted) { (void)api->functions().stream_synchronize(stream.get()); - poisoned = true; - position = 0; - saved.reset(); + if (!retry) { + poisoned = true; + position = 0; + saved.reset(); + } } throw; } } + /// \brief one prompt, one page at a time when the cache is paged + /// \note Starts at `position`, which is 0 or a page boundary when `ids` + /// holds more than one row. + buffer RunPrompt(std::span ids) { + if (!paged || ids.size() <= 1 || + (position == 0 && static_cast(ids.size()) <= kPageSize)) + return Run(ids); + buffer out; + const auto step = static_cast(kPageSize); + for (std::size_t start = 0; start < ids.size(); start += step) { + const auto count = std::min(step, ids.size() - start); + const auto chunk = ids.subspan(start, count); + if (position == 0) { + out = Run(chunk); + continue; + } + for (int attempt = 1;; ++attempt) { + try { + out = Run(chunk); + break; + } catch (const std::runtime_error& error) { + const bool again = + std::string(error.what()).find("non-finite") != std::string::npos && + attempt < kPagedChunkAttempts; + if (!again) throw; + std::cout << "[FLM] Gemma 4 paged prefill at position " << position + << " came back non-finite; running it again (" << attempt + << " of " << kPagedChunkAttempts << ")" << std::endl; + } + } + } + return out; + } + /// \brief refuse a non-finite result, naming where it came from /// \note CHECKED AT THE PASS'S OWN EXIT AND AGAIN AT THE HEAD, because /// both are places a caller reads from. Without it an all-NaN result @@ -1630,7 +1731,8 @@ struct gemma4_rai::Impl { buffer ReadCache(bool is_k, int layer, int index) { usable(); const auto& cfg = shape; - if (layer < 0 || layer >= cfg.layers || index < 0 || index >= kMaxSequenceLength) + if (layer < 0 || layer >= cfg.layers || index < 0 || + index >= static_cast(cache_rows)) throw std::out_of_range("Gemma 4 cache index is out of range"); // THE CACHE A LAYER READS IS ITS OWNER'S, not its own index. Twenty of // E2B's layers own none at all. @@ -1643,8 +1745,7 @@ struct gemma4_rai::Impl { buffer out(static_cast(cfg.kv_heads * head)); void* cache = (is_k ? k_cache : v_cache)[owner].get(); for (std::int64_t kv_head = 0; kv_head < cfg.kv_heads; ++kv_head) { - const auto offset = static_cast( - (kv_head * kMaxSequenceLength + index) * head); + const auto offset = CacheOffset(kv_head, index, head); api->Check(api->functions().tensor_read( cache, ryzenai_corelib_data_type_bf16, out.data() + kv_head * head, static_cast(head), @@ -1692,23 +1793,27 @@ buffer gemma4_rai::prefill(std::vector& ids, void*) { auto& impl = *impl_; if (impl.position == 0 || ids.size() <= 1) { if (impl.position == 0) impl.history.clear(); - auto out = impl.Run(std::span(ids.data(), ids.size())); + auto out = impl.RunPrompt(std::span(ids.data(), ids.size())); impl.history.insert(impl.history.end(), ids.begin(), ids.end()); return out; } - // A multi-row continuation re-runs the whole conversation from position 0. + // A multi-row pass starts at position 0 or, on a paged cache, on a page + // boundary. So a continuation re-runs the rows since the nearest such + // start ahead of its own: the whole conversation unpaged, and on a paged + // cache only the current page's rows, none at all on a boundary. impl.usable(); if (impl.history.size() != static_cast(impl.position)) throw std::invalid_argument("Gemma 4 cannot continue from position " + std::to_string(impl.position) + " without the tokens behind it"); - std::vector tokens = impl.history; - tokens.insert(tokens.end(), ids.begin(), ids.end()); const int resume = impl.position; - impl.position = 0; + const int from = impl.paged ? resume - resume % static_cast(kPageSize) : 0; + std::vector tokens(impl.history.begin() + from, impl.history.end()); + tokens.insert(tokens.end(), ids.begin(), ids.end()); + impl.position = from; try { - auto out = impl.Run(std::span(tokens.data(), tokens.size())); - impl.history = std::move(tokens); + auto out = impl.RunPrompt(std::span(tokens.data(), tokens.size())); + impl.history.insert(impl.history.end(), ids.begin(), ids.end()); return out; } catch (...) { if (!impl.poisoned) impl.position = resume; @@ -1739,7 +1844,7 @@ void gemma4_rai::load_weights(Q4NX&) { void gemma4_rai::update_max_length(std::uint32_t max_length) { impl_->usable(); - if (!max_length || max_length > kMaxSequenceLength || + if (!max_length || max_length > impl_->cache_rows || max_length < static_cast(impl_->position)) throw std::out_of_range("Gemma 4 maximum length is invalid"); impl_->max_length = max_length; diff --git a/src/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp b/src/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp index 72aea47fe..7d90cfc50 100644 --- a/src/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp +++ b/src/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp @@ -132,11 +132,12 @@ class RaiBackend final : public flm::backend::ModelBackend { throw std::invalid_argument( "Gemma 4 rai does not support preemption"); } - if (context.context_length < 1 || - context.context_length > kRaiContextLimit) { + if (!ContextFits(context.context_length)) { throw std::out_of_range( - "Gemma 4 rai context length must be in 1..4096"); + "Gemma 4 rai context length must be in 1..4096, or a multiple of 4096 up to " + + std::to_string(kMaxPagedContext)); } + context_length_ = context.context_length; // 2-3. Read and cross-validate EVERY source of truth before the // runtime is acquired or the engine built, so a mismatched @@ -201,7 +202,9 @@ class RaiBackend final : public flm::backend::ModelBackend { : std::string(); } - std::uint32_t max_decode_length() const override { return kRaiDecodeLimit; } + std::uint32_t max_decode_length() const override { + return context_length_ > 0 ? context_length_ - 1 : 0; + } bool supports_preemption() const override { return false; } @@ -222,6 +225,7 @@ class RaiBackend final : public flm::backend::ModelBackend { private: /// \brief read from the package at construction; see the constructor + std::uint32_t context_length_{}; std::vector forced_eos_ids_; /// \brief read from the package at construction; see the constructor int forced_bos_id_{}; diff --git a/src/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp b/src/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp index aa0eb8478..f73390736 100644 --- a/src/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp +++ b/src/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp @@ -288,6 +288,11 @@ Gemma4Config ConfigFromMetadata(const flm::rai::GgufFile& file) { const auto sliding_window = static_cast(file.Unsigned("gemma4.attention.sliding_window")); + // A file that does not state its context gets one page, as corelib's + // driver does. + const auto context_length = file.HasMetadata("gemma4.context_length") + ? static_cast(file.Unsigned("gemma4.context_length")) + : kMaxSequenceLength; const auto eps = static_cast(file.Number("gemma4.attention.layer_norm_rms_epsilon")); const auto rope_theta = file.Number("gemma4.rope.freq_base"); @@ -361,6 +366,7 @@ Gemma4Config ConfigFromMetadata(const flm::rai::GgufFile& file) { config.group = row.group; config.head_group = row.head_group; config.ple_group = row.ple_group; + config.context_length = context_length; return config; } diff --git a/src/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp b/src/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp index caa584b8f..9b7c99f5a 100644 --- a/src/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp +++ b/src/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp @@ -138,11 +138,27 @@ ryzenai_corelib_flat_mha_bf16_desc BuildAttentionDesc(const Gemma4Config& config return desc; } +ryzenai_corelib_paged_mha_bf16_desc AsPaged(const ryzenai_corelib_flat_mha_bf16_desc& flat, + std::int64_t context) { + ryzenai_corelib_paged_mha_bf16_desc desc{}; + desc.num_heads = flat.num_heads; + desc.kv_num_heads = flat.kv_num_heads; + desc.head_size = flat.head_size; + desc.max_seq = context; + desc.page_size = kPageSize; + desc.rope_dim = flat.rope_dim; + desc.window = flat.window; + desc.kv_shared = flat.kv_shared; + desc.scale = flat.scale; + return desc; +} + } // namespace Gemma4ShapePlan Gemma4ShapePlan::Build( const std::shared_ptr& api, - ryzenai_corelib_stream_ptr stream, const Gemma4Config& config) { + ryzenai_corelib_stream_ptr stream, const Gemma4Config& config, + std::int64_t context) { if (!api) throw std::invalid_argument("Gemma4ShapePlan corelib API is null"); if (!stream) throw std::invalid_argument("Gemma4ShapePlan stream is null"); @@ -155,6 +171,17 @@ Gemma4ShapePlan Gemma4ShapePlan::Build( BuildAttentionDesc(config, swa, shared); } } + plan.paged_ = context > kPageSize; + if (plan.paged_) { + if (context > kMaxPagedContext || context % kPageSize != 0) { + throw std::invalid_argument( + "Gemma 4 paged context must be a multiple of 4096 up to " + + std::to_string(kMaxPagedContext)); + } + for (std::size_t index = 0; index < plan.attention_descs_.size(); ++index) { + plan.paged_descs_[index] = AsPaged(plan.attention_descs_[index], context); + } + } // ssmlp must be planned at EVERY DISTINCT width in // config.layer_intermediate, not just the first layer's -- E2B steps @@ -175,10 +202,26 @@ Gemma4ShapePlan Gemma4ShapePlan::Build( plan.full_rows_.reserve(static_cast(kMaxSequenceLength)); const auto matmul = corelib::ShapeGrid::Matmul(*api, stream); const auto ssmlp_grid = corelib::ShapeGrid::SsMlp(*api, stream, true); - const auto sliding_mha = corelib::ShapeGrid::FlatMha( - *api, stream, plan.attention_descs_[DescIndex(true, false)]); - const auto full_mha = corelib::ShapeGrid::FlatMha( - *api, stream, plan.attention_descs_[DescIndex(false, false)]); + // A multi-row pass at position 0 is attention over the first page, and + // corelib's paged_mha runs it on flat_mha's own prefill kernel at the + // pass's rows; every later chunk is a whole page on the paged kernel. So + // the row extents come from flat_mha's grid even on a paged cache, and + // the paged grid only has to carry the one page-sized chunk. + const auto sliding_mha = + corelib::ShapeGrid::FlatMha(*api, stream, plan.attention_descs_[DescIndex(true, false)]); + const auto full_mha = + corelib::ShapeGrid::FlatMha(*api, stream, plan.attention_descs_[DescIndex(false, false)]); + if (plan.paged_) { + for (const bool swa : {true, false}) { + const auto paged = corelib::ShapeGrid::PagedMha( + *api, stream, plan.paged_descs_[DescIndex(swa, false)]); + const char* what = swa ? "sliding paged_mha chunk" : "full paged_mha chunk"; + if (corelib::CoveringRows(paged, kPageSize, 0, 0, -1, what) != kPageSize) { + throw std::runtime_error(std::string(what) + " has no kernel at exactly " + + std::to_string(kPageSize) + " rows"); + } + } + } for (const auto rows : kExecutionRows) { // ssmlp: gelu (activation = 1) with the post-feedforward norm SET @@ -321,6 +364,11 @@ Gemma4ShapePlan::attention_desc(bool swa, bool shared) const noexcept { return attention_descs_[DescIndex(swa, shared)]; } +const ryzenai_corelib_paged_mha_bf16_desc& +Gemma4ShapePlan::paged_attention_desc(bool swa, bool shared) const noexcept { + return paged_descs_[DescIndex(swa, shared)]; +} + std::int64_t Gemma4ShapePlan::lm_head_rows() const noexcept { return lm_head_rows_; } } // namespace flm::gemma4 diff --git a/src/common/models/phi4/rai/aie_next/phi4_rai.cpp b/src/common/models/phi4/rai/aie_next/phi4_rai.cpp index e99adf829..48a46d4e5 100644 --- a/src/common/models/phi4/rai/aie_next/phi4_rai.cpp +++ b/src/common/models/phi4/rai/aie_next/phi4_rai.cpp @@ -138,8 +138,8 @@ struct phi4_rai::Impl { first_norm_scale = an[0]; phases.host_prep = phases.Lap(); - auto lease=runtime->AcquireExecution(); void* raw=nullptr; - api->Check(api->functions().create_stream(kPrefillPdi,kTokenPdi,&raw),"ryzenai_corelib_create_stream"); stream=UniqueStream(api,raw); + auto lease=runtime->AcquireExecution(); + stream=corelib::OpenModelStream(api,kPrefillPdi,kTokenPdi,"phi4mini"); // Now that the stream exists the padding helpers can be asked. Every // extent below is therefore an extent *for this stream's PDI pair*. plan=Phi4ShapePlan::Build(api,stream.get()); @@ -426,7 +426,7 @@ struct phi4_rai::Impl { const auto& extents=plan.ForRows(ids.size()); auto view=[&](void* parent,std::int64_t rows,std::int64_t cols,const char* label){ std::array shape{rows,cols};void* p=nullptr; - api->Check(api->functions().create_tensor_window(parent,shape.data(),shape.size(),0,&p),std::string("ryzenai_corelib_create_tensor_window ")+label); + api->Check(api->functions().create_tensor_window(parent,shape.data(),shape.size(),0,nullptr,0,&p),std::string("ryzenai_corelib_create_tensor_window ")+label); return UniqueTensorWindow(api,p); }; const auto hidden_q=view(hidden.get(),extents.query_rows,kHiddenSize,"hidden@query"); @@ -452,8 +452,8 @@ struct phi4_rai::Impl { // The gm exception: a cache write takes its position from the // bound buffer's base, so this window stays a position, not an // extent -- its height is what bounds the rows that land there. - std::array shape{8,kMaxSequenceLength-position,128};void* p=nullptr; - api->Check(api->functions().create_tensor_window(v_cache[i].get(),shape.data(),shape.size(),static_cast(position)*128,&p),"ryzenai_corelib_create_tensor_window V");UniqueTensorWindow win(api,p); + std::array shape{8,kMaxSequenceLength-position,128};const std::array strides{kMaxSequenceLength*128,128,1};void* p=nullptr; + api->Check(api->functions().create_tensor_window(v_cache[i].get(),shape.data(),shape.size(),static_cast(position)*128,strides.data(),strides.size(),&p),"ryzenai_corelib_create_tensor_window V");UniqueTensorWindow win(api,p); api->Check(api->functions().matmul(stream.get(),hidden_kv.get(),v_weights[i].get(),win.get()),"ryzenai_corelib_matmul_bf16 value layer "+std::to_string(i)); api->Check(api->functions().flat_mha(stream.get(),&plan.attention_desc(),q_mha.get(),k_mha.get(),position,cosine.get(),sine.get(),k_cache[i].get(),v_cache[i].get(),attention_mha.get()),"ryzenai_corelib_flat_mha_bf16 layer "+std::to_string(i)); api->Check(api->functions().matmul(stream.get(),attention_mm.get(),o_weights[i].get(),hidden_out.get()),"ryzenai_corelib_matmul_bf16 output layer "+std::to_string(i)); diff --git a/src/common/models/qwen3/rai/aie_next/qwen3_rai.cpp b/src/common/models/qwen3/rai/aie_next/qwen3_rai.cpp index 8e63d5358..23810a576 100644 --- a/src/common/models/qwen3/rai/aie_next/qwen3_rai.cpp +++ b/src/common/models/qwen3/rai/aie_next/qwen3_rai.cpp @@ -123,10 +123,7 @@ struct qwen3_rai::Impl { phases.host_prep = phases.Lap(); auto lease = runtime->AcquireExecution(); - void* raw = nullptr; - api->Check(api->functions().create_stream(kPrefillPdi, kTokenPdi, &raw), - "ryzenai_corelib_create_stream"); - stream = UniqueStream(api, raw); + stream = corelib::OpenModelStream(api, kPrefillPdi, kTokenPdi, "qwen3"); plan = Qwen3ShapePlan::Build(api, stream.get(), *c); phases.shape_plan = phases.Lap(); @@ -465,7 +462,7 @@ struct qwen3_rai::Impl { const auto view = [&](void* parent, std::int64_t r, std::int64_t cols, const char* label) { const std::array shape{r, cols}; void* p = nullptr; - api->Check(f.create_tensor_window(parent, shape.data(), shape.size(), 0, &p), + api->Check(f.create_tensor_window(parent, shape.data(), shape.size(), 0, nullptr, 0, &p), std::string("ryzenai_corelib_create_tensor_window ") + label); return UniqueTensorWindow(api, p); }; @@ -498,9 +495,12 @@ struct qwen3_rai::Impl { // this pass's first row, whose height bounds what lands there. const std::array v_shape{c->kv_heads, kMaxSequenceLength - position, c->head_dim}; + const std::array v_strides{kMaxSequenceLength * c->head_dim, + c->head_dim, 1}; void* p = nullptr; api->Check(f.create_tensor_window(v_cache[i].get(), v_shape.data(), v_shape.size(), - static_cast(position * c->head_dim), &p), + static_cast(position * c->head_dim), + v_strides.data(), v_strides.size(), &p), "ryzenai_corelib_create_tensor_window V"); const UniqueTensorWindow v_window(api, p); api->Check(f.matmul(stream.get(), hidden_w.get(), v_weights[i].get(), v_window.get()), diff --git a/src/common/models/qwen35/rai/aie_next/qwen35_rai.cpp b/src/common/models/qwen35/rai/aie_next/qwen35_rai.cpp index 04431da44..ba413319b 100644 --- a/src/common/models/qwen35/rai/aie_next/qwen35_rai.cpp +++ b/src/common/models/qwen35/rai/aie_next/qwen35_rai.cpp @@ -122,15 +122,7 @@ class qwen35_rai::Impl { const int layers = cfg->layers; auto lease = this->runtime->AcquireExecution(); - void* raw = nullptr; - api->Check(api->functions().create_stream(kPrefillPdi, kTokenPdi, &raw), - "ryzenai_corelib_create_stream"); - stream = UniqueStream(api, raw); - const char* root = api->functions().stream_get_kernels_root(stream.get()); - if (!root || !*root) - throw std::runtime_error( - "this stream has no kernel root. Set RYZENAI_CORELIB_DD_ROOT to the existing " - "DynamicDispatch checkout; corelib reads kernels from its transaction/mds"); + stream = corelib::OpenModelStream(api, kPrefillPdi, kTokenPdi, "qwen35"); const auto matmul = ShapeGrid::Matmul(*api, stream.get()); const auto sigmoid = ShapeGrid::MatmulAct( @@ -269,10 +261,13 @@ class qwen35_rai::Impl { } UniqueTensor Window(void* parent, std::initializer_list shape, std::size_t offset, - const char* label) { + const char* label, std::initializer_list strides = {}) { std::vector dims(shape); + std::vector steps(strides); void* created = nullptr; - api->Check(api->functions().create_tensor_window(parent, dims.data(), dims.size(), offset, &created), + api->Check(api->functions().create_tensor_window(parent, dims.data(), dims.size(), offset, + steps.empty() ? nullptr : steps.data(), + steps.size(), &created), label); return UniqueTensor(api, created); } @@ -585,7 +580,7 @@ class qwen35_rai::Impl { io.dw_q.get(), io.dw_k.get(), io.dw_v.get(), io.conv[1 - side].get()), "dwconv"); api->Check(f.linear_attention(stream.get(), &lin_desc, io.state[side].get(), io.dw_q.get(), io.dw_k.get(), - io.dw_v.get(), io.g_live.get(), io.beta_live.get(), io.lin_out.get(), + io.dw_v.get(), io.g_live.get(), io.beta_live.get(), 1, io.lin_out.get(), io.state[1 - side].get()), "linear attention"); api->Check(f.rmsnorm(stream.get(), io.lin_out.get(), weights.norm.get(), io.lin_out.get()), @@ -604,7 +599,8 @@ class qwen35_rai::Impl { auto& v_row = io.v_rows[static_cast(position)]; if (!v_row) { v_row = Window(io.v_cache.get(), {cfg->kv_heads, 1, cfg->head_dim}, - static_cast(position) * cfg->head_dim, "v row"); + static_cast(position) * cfg->head_dim, "v row", + {kMaxSequenceLength * cfg->head_dim, cfg->head_dim, 1}); } api->Check(f.matmul(stream.get(), hidden, weights.q.get(), io.q_raw.get()), "q"); api->Check(f.matmul(stream.get(), hidden, weights.k.get(), io.k_live.get()), "k"); diff --git a/src/common/models/qwen36/rai/aie_next/qwen36_rai.cpp b/src/common/models/qwen36/rai/aie_next/qwen36_rai.cpp new file mode 100644 index 000000000..caef9bc8a --- /dev/null +++ b/src/common/models/qwen36/rai/aie_next/qwen36_rai.cpp @@ -0,0 +1,289 @@ +#include "models/qwen36/rai/aie_next/qwen36_rai.hpp" + +#include "rai/corelib_object.hpp" +#include "rai/host_ops.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace flm::qwen36 { +namespace { + +using flm::corelib::CorelibApi; +using flm::corelib::CorelibRuntime; +using flm::corelib::UniqueModel; +using flm::corelib::UniqueModelSnapshot; +using flm::corelib::UniqueModelWeights; + +} // namespace + +class qwen36_rai::Impl { +public: + Impl(std::filesystem::path gguf, std::shared_ptr runtime, + std::uint32_t maximum) + : runtime(std::move(runtime)), + api(this->runtime ? this->runtime->api() : nullptr), + max_length(maximum) { + if (gguf.empty()) throw std::invalid_argument("Qwen3.6 GGUF path is empty"); + if (!this->runtime || !api) throw std::invalid_argument("corelib runtime is null"); + if (!maximum || maximum > kContextLimit) + throw std::invalid_argument("Qwen3.6 maximum length must be in 1.." + + std::to_string(kContextLimit)); + + std::cout << "[FLM] Loading Qwen3.6-35B-A3B from Q8_0 GGUF" << std::endl; + auto lease = this->runtime->AcquireExecution(); + const auto path = gguf.string(); + void* raw_weights = nullptr; + api->Check(api->functions().model_weights_create_path("qwen36", path.c_str(), &raw_weights), + "ryzenai_corelib_model_weights_create_path"); + weights = UniqueModelWeights(api, raw_weights); + Recreate(maximum); + logits.resize(vocab); + } + + buffer Forward(int id) { + usable(); + auto lease = runtime->AcquireExecution(); + Run(&id, 1, true); + history.push_back(id); + return Logits(); + } + + buffer Prefill(const std::vector& ids, void* payload) { + usable(); + if (payload) + throw std::invalid_argument("Qwen3.6 rai does not run a vision payload"); + if (ids.empty()) throw std::invalid_argument("Qwen3.6 request contains no token IDs"); + auto lease = runtime->AcquireExecution(); + Run(ids.data(), ids.size(), true); + history.insert(history.end(), ids.begin(), ids.end()); + return Logits(); + } + + void Clear() { + usable(); + auto lease = runtime->AcquireExecution(); + ResetModel(); + history.clear(); + saved.reset(); + } + + void Rewind(int to) { + usable(); + if (to < 0 || static_cast(to) > max_length) + throw std::out_of_range("Qwen3.6 context length is out of range"); + if (static_cast(to) > history.size()) + throw std::runtime_error("Qwen3.6 cannot rewind past the tokens it has run"); + if (to == position) return; + auto lease = runtime->AcquireExecution(); + if (saved && saved->position == to && history.size() >= saved->history.size() && + std::equal(saved->history.begin(), saved->history.end(), history.begin())) { + Load(saved->snapshot, to); + history = saved->history; + return; + } + history.resize(static_cast(to)); + ResetModel(); + if (to > 0) Run(history.data(), history.size(), false); + } + + int Checkpoint() { + usable(); + if (Matches(saved)) return position; + auto lease = runtime->AcquireExecution(); + saved = Saved{position, history, Save()}; + return position; + } + + int Restore() { + usable(); + if (!saved) return -1; + if (Matches(saved)) return position; + auto lease = runtime->AcquireExecution(); + Load(saved->snapshot, saved->position); + history = saved->history; + return position; + } + + void Grow(std::uint32_t n) { + if (n <= capacity) return; + auto lease = runtime->AcquireExecution(); + UniqueModelSnapshot current; + if (position > 0) current = Save(); + Recreate(n); + if (current) Load(current, position); + } + + void usable() const { + if (poisoned) throw std::runtime_error("Qwen3.6 corelib engine is poisoned"); + } + + struct Saved { + int position = 0; + std::vector history; + UniqueModelSnapshot snapshot; + }; + + bool poisoned = false; + + std::shared_ptr runtime; + std::shared_ptr api; + UniqueModelWeights weights; + UniqueModel model; + std::uint32_t capacity = 0; + std::uint32_t max_length = 0; + std::size_t vocab = 0; + int position = 0; + std::vector history; + std::optional saved; + std::vector logits; + +private: + /// \brief a model over `weights` holding at least `n` positions + /// \note The capacity is fixed for a model's life; `weights` outlive it, so + /// a larger one is a new model over the same packed weights. + void Recreate(std::uint32_t n) { + ryzenai_corelib_model_create_options options{}; + options.size = sizeof(options); + options.max_sequence_length = static_cast(n); + void* raw = nullptr; + api->Check(api->functions().model_create("qwen36", weights.get(), &options, &raw), + "ryzenai_corelib_model_create"); + UniqueModel created(api, raw); + ryzenai_corelib_model_info info{}; + api->Check(api->functions().model_info_get(created.get(), &info), + "ryzenai_corelib_model_info_get"); + if (info.vocab_size <= 0 || info.context_length < static_cast(n)) { + throw std::runtime_error( + "Qwen3.6 model_info is vocab " + std::to_string(info.vocab_size) + + " context " + std::to_string(info.context_length) + + "; this engine needs vocab > 0 and context >= " + std::to_string(n)); + } + model = std::move(created); + capacity = static_cast(info.context_length); + vocab = static_cast(info.vocab_size); + } + + void Run(const int* ids, std::size_t rows, bool want_logits) { + if (static_cast(position) + rows > max_length) + throw std::out_of_range("Qwen3.6 request exceeds configured context capacity"); + const std::vector tokens(ids, ids + rows); + ryzenai_corelib_model_forward_options options{}; + options.position = RYZENAI_CORELIB_POSITION_CONTINUE; + if (want_logits) { + options.logits = logits.data(); + options.logits_capacity = logits.size(); + } + try { + api->Check(api->functions().model_forward( + model.get(), tokens.data(), nullptr, rows, nullptr, &options), + "ryzenai_corelib_model_forward"); + } catch (...) { + poisoned = true; + position = 0; + history.clear(); + saved.reset(); + throw; + } + position += static_cast(rows); + } + + UniqueModelSnapshot Save() { + void* raw = nullptr; + api->Check(api->functions().model_state_save(model.get(), &raw), + "ryzenai_corelib_model_state_save"); + return UniqueModelSnapshot(api, raw); + } + + /// \note corelib leaves the model reset when a restore fails partway, so a + /// failure leaves nothing this engine's position or history describe. + void Load(const UniqueModelSnapshot& snapshot, int at) { + try { + api->Check(api->functions().model_state_restore(model.get(), snapshot.get()), + "ryzenai_corelib_model_state_restore"); + } catch (...) { + poisoned = true; + position = 0; + history.clear(); + saved.reset(); + throw; + } + position = at; + } + + bool Matches(const std::optional& snapshot) const { + return snapshot && snapshot->position == position && snapshot->history == history; + } + + buffer Logits() const { + buffer out(vocab); + rai::F32ToBf16(logits, std::span( + reinterpret_cast(out.data()), vocab)); + return out; + } + + void ResetModel() { + api->Check(api->functions().model_reset(model.get()), "ryzenai_corelib_model_reset"); + position = 0; + } +}; + +qwen36_rai::qwen36_rai(std::filesystem::path gguf, + std::shared_ptr runtime, + std::uint32_t max_length) + : impl_(std::make_unique(std::move(gguf), std::move(runtime), max_length)) {} + +qwen36_rai::~qwen36_rai() = default; + +buffer qwen36_rai::forward(int id) { return impl_->Forward(id); } + +buffer qwen36_rai::prefill(std::vector& ids, void* payload) { + return impl_->Prefill(ids, payload); +} + +void qwen36_rai::set_context_length(int n) { impl_->Rewind(n); } + +void qwen36_rai::load_weights(Q4NX&) { + impl_->usable(); + throw std::runtime_error("Qwen3.6 rai weights are loaded only from GGUF"); +} + +void qwen36_rai::update_max_length(std::uint32_t n) { + impl_->usable(); + if (!n || n > kContextLimit || n < static_cast(impl_->position)) + throw std::out_of_range("Qwen3.6 maximum length is invalid"); + impl_->Grow(n); + if (impl_->saved && n < static_cast(impl_->saved->position)) + impl_->saved.reset(); + impl_->max_length = n; +} + +void qwen36_rai::clear_context() { impl_->Clear(); } + +buffer qwen36_rai::get_k_cache(int, int) { + throw std::runtime_error("Qwen3.6 rai does not expose the KV cache"); +} + +buffer qwen36_rai::get_v_cache(int, int) { + throw std::runtime_error("Qwen3.6 rai does not expose the KV cache"); +} + +int qwen36_rai::get_current_context_length() { + impl_->usable(); + return impl_->position; +} + +int qwen36_rai::checkpoint() { return impl_->Checkpoint(); } + +int qwen36_rai::restore() { return impl_->Restore(); } + +bool qwen36_rai::poisoned() const noexcept { return impl_ && impl_->poisoned; } + +} // namespace flm::qwen36 diff --git a/src/common/models/qwen36/rai/aie_next/qwen36_rai_backend.cpp b/src/common/models/qwen36/rai/aie_next/qwen36_rai_backend.cpp new file mode 100644 index 000000000..23d1138f4 --- /dev/null +++ b/src/common/models/qwen36/rai/aie_next/qwen36_rai_backend.cpp @@ -0,0 +1,109 @@ +#include "models/qwen36/rai/aie_next/qwen36_rai_backend.hpp" + +#include "models/qwen36/rai/aie_next/qwen36_rai.hpp" +#include "rai/corelib_runtime.hpp" +#include "utils/utils.hpp" + +#include +#include +#include +#include +#include +#include +#include + +namespace flm::qwen36 { +namespace { + +std::string GgufName(const nlohmann::ordered_json& model_info) { + std::string found; + if (model_info.contains("files") && model_info["files"].is_array()) { + for (const auto& file : model_info["files"]) { + if (!file.is_string()) continue; + const auto name = file.get(); + if (name.size() > 5 && name.ends_with(".gguf")) { + if (!found.empty()) + throw std::runtime_error("Qwen3.6 rai entry lists more than one .gguf file"); + found = name; + } + } + } + if (found.empty()) throw std::runtime_error("Qwen3.6 rai entry lists no .gguf file"); + return found; +} + +// The published tokenizer_config names eos_token as a string and does not +// carry eos_token_id. The id is the added token of that same string. +std::vector EosIds(const std::filesystem::path& model) { + std::ifstream config_file(model / "tokenizer_config.json"); + if (!config_file) throw std::runtime_error("Qwen3.6 tokenizer_config.json is missing"); + const auto config = nlohmann::json::parse(config_file); + if (!config.contains("eos_token") || !config["eos_token"].is_string()) + throw std::runtime_error("Qwen3.6 tokenizer_config.json has no eos_token string"); + const auto eos = config["eos_token"].get(); + + std::ifstream tokenizer_file(model / "tokenizer.json"); + if (!tokenizer_file) throw std::runtime_error("Qwen3.6 tokenizer.json is missing"); + const auto tokenizer = nlohmann::json::parse(tokenizer_file); + std::vector ids; + if (tokenizer.contains("added_tokens") && tokenizer["added_tokens"].is_array()) { + for (const auto& token : tokenizer["added_tokens"]) { + if (token.value("content", std::string()) == eos && token.contains("id") && + token["id"].is_number_integer()) + ids.push_back(token["id"].get()); + } + } + if (ids.empty()) + throw std::runtime_error("Qwen3.6 tokenizer.json has no id for eos token " + eos); + return ids; +} + +class RaiBackend final : public flm::backend::ModelBackend { +public: + explicit RaiBackend(const flm::backend::BackendContext& context) { + if (context.enable_preemption) + throw std::invalid_argument("Qwen3.6 rai does not support preemption"); + if (context.context_length < 1 || context.context_length > kRaiContextLimit) + throw std::out_of_range("Qwen3.6 rai context length must be in 1.." + + std::to_string(kRaiContextLimit)); + context_length_ = context.context_length; + eos_ids_ = EosIds(context.model_path); + + runtime_ = corelib::CorelibRuntime::GetOrCreate(utils::get_executable_directory()); + auto engine = std::make_unique( + std::filesystem::path(context.model_path) / GgufName(context.model_info), + runtime_, context.context_length); + engine->clear_context(); + engine_ = std::move(engine); + } + + ~RaiBackend() override { engine_.reset(); } + + causal_lm& engine() override { return *engine_; } + std::string id() const override { return flm::backend::kRaiBackendId; } + std::string detail() const override { + return runtime_ ? runtime_->loaded_library_path().string() : std::string(); + } + std::uint32_t max_decode_length() const override { return context_length_ - 1; } + bool supports_preemption() const override { return false; } + bool forwards_past_eos() const override { return false; } + bool poisoned() const noexcept override { return engine_ && engine_->poisoned(); } + std::optional> forced_eos_ids() const override { return eos_ids_; } + +private: + std::uint32_t context_length_{}; + std::vector eos_ids_; + std::shared_ptr runtime_; + std::unique_ptr engine_; +}; + +} // namespace + +flm::backend::BackendFactory rai_factory() { + return [](const flm::backend::BackendContext& context) + -> std::unique_ptr { + return std::make_unique(context); + }; +} + +} // namespace flm::qwen36 diff --git a/src/common/modules/sampler.cpp b/src/common/modules/sampler.cpp index 50b778ef1..bd255481f 100644 --- a/src/common/modules/sampler.cpp +++ b/src/common/modules/sampler.cpp @@ -215,33 +215,38 @@ void Sampler::sampler_penalty_apply_sparse() { } void Sampler::sampler_topk_apply(int k) { - if (k == 0) { + k = std::min(k, in_features); + if (k <= 0) { return; } - logits_list_t pairs; - pairs.resize(in_features); - for (int i = 0; i < in_features; i++) { - pairs[i].logits = this->logits[i]; - pairs[i].token_id = i; - pairs[i].prob = 0.0f; - } - - std::partial_sort( - pairs.begin(), - pairs.begin() + k, - pairs.end(), - [](const logits_t& a, const logits_t& b) { - // bf16 logits tie exactly often enough to matter over a 100k+ vocab, - // and partial_sort is not stable, so a bare `>` leaves the winner of a - // tie down to the sort's internals. Lowest id wins instead: the order - // is then reproducible and matches sample_greedy()'s argmax. - if (a.logits != b.logits) { return a.logits > b.logits; } - return a.token_id < b.token_id; + // bf16 logits tie exactly often enough to matter over a 100k+ vocab, so a + // bare `>` would leave the winner of a tie to the selection's internals. + // Lowest id wins instead: the order is reproducible and matches + // sample_greedy()'s argmax. + const auto ranks_before = [](const logits_t& a, const logits_t& b) { + if (a.logits != b.logits) { return a.logits > b.logits; } + return a.token_id < b.token_id; + }; + + // One pass with a k-entry heap whose front is the worst survivor, instead of + // materialising and partially sorting a vocab-sized array every token. + auto& top = this->top_k_logits; + top.clear(); + top.reserve(k); + for (int i = 0; i < k; i++) { + top.push_back(logits_t{ this->logits[i], i, 0.0f }); + } + std::make_heap(top.begin(), top.end(), ranks_before); + for (int i = k; i < in_features; i++) { + // ids only grow, so a logit equal to the worst survivor's never displaces it + if (this->logits[i] > top.front().logits) { + std::pop_heap(top.begin(), top.end(), ranks_before); + top.back() = logits_t{ this->logits[i], i, 0.0f }; + std::push_heap(top.begin(), top.end(), ranks_before); } - ); - - this->top_k_logits.assign(pairs.begin(), pairs.begin() + k); + } + std::sort_heap(top.begin(), top.end(), ranks_before); } void Sampler::sampler_topp_apply(float p) { diff --git a/src/common/rai/corelib_api.cpp b/src/common/rai/corelib_api.cpp index a438eb048..1a447d2da 100644 --- a/src/common/rai/corelib_api.cpp +++ b/src/common/rai/corelib_api.cpp @@ -170,13 +170,54 @@ std::shared_ptr CorelibApi::Load(const std::filesystem::path& librar return ResolveForTest(std::move(resolver), absolute_library); } +#ifdef _WIN32 +namespace { +std::filesystem::path ModuleFileName(HMODULE module) { + if (!module) return {}; + // Long paths exist; grow until the name fits rather than truncate silently. + std::wstring buffer(MAX_PATH, L'\0'); + for (;;) { + const DWORD written = GetModuleFileNameW( + module, buffer.data(), static_cast(buffer.size())); + if (written == 0) return {}; + if (written < buffer.size()) { + buffer.resize(written); + return std::filesystem::path(buffer); + } + buffer.resize(buffer.size() * 2); + } +} +} // namespace +#endif + +std::filesystem::path CorelibApi::LoadedModulePath(const wchar_t* name) { +#ifdef _WIN32 + if (!name || !*name) return {}; + return ModuleFileName(GetModuleHandleW(name)); +#else + (void)name; + return {}; +#endif +} + #if defined(FLM_CORELIB_LINK_STATIC) std::shared_ptr CorelibApi::LoadStatic() { + // An import library still leaves a DLL bound by the loader's search order; + // report which copy that was. Asked by name: without dllimport, a linked + // symbol's address is the import thunk inside this executable. + const std::filesystem::path linked_from = LoadedModulePath(L"ryzenai_corelib.dll"); + const bool separate_dll = !linked_from.empty(); const char* configured = std::getenv("FLM_RAI_CORELIB_PATH"); if (configured && *configured) { - std::cerr << "[FLM] Ignoring FLM_RAI_CORELIB_PATH: corelib is linked " - "into this build, so there is no DLL to select." - << std::endl; + std::cerr << "[FLM] Ignoring FLM_RAI_CORELIB_PATH: this build binds " + "corelib at link time" + << (separate_dll ? ", and the Windows loader has already " + "loaded it from '" + linked_from.string() + "'" + : "") + << "." << std::endl; + } + if (separate_dll) { + std::cerr << "[FLM] corelib: " << linked_from.string() << std::endl; } // Resolve through the same name-keyed resolver the DLL path uses, so the // version check and the missing-symbol diagnostics stay identical. @@ -190,7 +231,8 @@ std::shared_ptr CorelibApi::LoadStatic() { return nullptr; }; return ResolveForTest(std::move(resolver), - std::filesystem::path("")); + separate_dll ? linked_from + : std::filesystem::path("")); } #endif diff --git a/src/common/rai/host_ops.cpp b/src/common/rai/host_ops.cpp index f4180a045..a08aacadb 100644 --- a/src/common/rai/host_ops.cpp +++ b/src/common/rai/host_ops.cpp @@ -97,18 +97,23 @@ void RmsNorm(std::span input, std::span scale, } std::vector F32ToBf16(std::span values) { - std::vector result; - result.reserve(values.size()); - for (const float value : values) { - std::uint32_t bits = std::bit_cast(value); + std::vector result(values.size()); + F32ToBf16(values, result); + return result; +} + +void F32ToBf16(std::span values, std::span output) { + if (output.size() != values.size()) + throw std::invalid_argument("F32ToBf16 output size does not match its input"); + for (std::size_t i = 0; i < values.size(); ++i) { + std::uint32_t bits = std::bit_cast(values[i]); if ((bits & 0x7fffffffU) > 0x7f800000U) { bits |= 0x00400000U; } else { bits += 0x7fffU + ((bits >> 16) & 1U); } - result.push_back(static_cast(bits >> 16)); + output[i] = static_cast(bits >> 16); } - return result; } } // namespace flm::rai diff --git a/src/home_install.sh b/src/home_install.sh index fca6889e9..b2c8ee896 100755 --- a/src/home_install.sh +++ b/src/home_install.sh @@ -13,8 +13,8 @@ # # --rai installs to /scratch/$USER/flm_exe_rai rather than ~/flm_exe: a rai # build and a stock one would otherwise overwrite each other, and the corelib -# and DynamicDispatch libraries staged into a rai prefix are far too large for a -# home directory. Falls back to ~/flm_exe_rai with no scratch space, and +# library staged into a rai prefix is far too large for a home directory. +# Falls back to ~/flm_exe_rai with no scratch space, and # FLM_PREFIX overrides either default. # # --rai selects the linux-rai-on preset. Point that build at a corelib prefix @@ -100,9 +100,8 @@ XRT_ROOT="" # and lets the two env scripts be sourced independently. # # That default lives on scratch rather than under $HOME: a rai install also -# carries libryzenai_corelib.so and the DynamicDispatch core it links against, -# which together dwarf the rest of the tree and do not belong in a home -# directory. $HOME is the fallback for a machine with no scratch space. +# carries libryzenai_corelib.so, whose built-in model kernels dwarf the rest of +# the tree and do not belong in a home directory. $HOME is the fallback for a machine with no scratch space. RAI_PREFIX_DEFAULT="/scratch/$USER/flm_exe_rai" [[ -d "/scratch/$USER" ]] || RAI_PREFIX_DEFAULT="$HOME/flm_exe_rai" if [[ -n "$FLM_PREFIX_FROM_ENV" ]]; then @@ -203,10 +202,10 @@ rm -f "$FLM_PREFIX/share/flm/model_list_rai.json" # ---- stage the rai (ryzenai-corelib) runtime ------------------------------- # A corelib-backed build links libryzenai_corelib.so, which no install rule # covers: CMake only globs lib/ for engine libraries, and corelib is -# not one of those. Copying it alone is still not enough -- it reaches -# libdyn_dispatch_core.so through an absolute RUNPATH into the out-of-tree -# DynamicDispatch prefix it was built against, so the install would break as -# soon as that directory moved. +# not one of those. Copying it alone is still not enough -- any non-system +# library it links is reached through an absolute RUNPATH into the prefix it +# was built against, so the install would break as soon as that directory +# moved. # # Copy corelib together with every non-system library it pulls in, leaving the # prefix self-contained. The dynamic loader searches LD_LIBRARY_PATH (set by the @@ -216,7 +215,7 @@ rm -f "$FLM_PREFIX/share/flm/model_list_rai.json" # Keyed on --rai rather than on the cache: a stock install has to be the stock # install main ships, and reading the flag out of whichever build tree happened # to be lying around made a plain ./home_install.sh drop libryzenai_corelib.so -# and a 400 MB libdyn_dispatch_core.so into ~/flm_exe, where the bundled XRT +# and its dependencies into ~/flm_exe, where the bundled XRT # then had company it could not cope with. A tree that disagrees with the mode # gets a warning instead. CACHE_FILE="$BUILD_DIR/CMakeCache.txt" @@ -247,36 +246,22 @@ elif [[ -f "$CACHE_FILE" ]] && grep -q '^FLM_ENABLE_RAI:BOOL=ON' "$CACHE_FILE"; install -m 0755 "$CORELIB" "$RAI_LIB_DEST/" # ldd resolves the whole transitive chain, so filtering out the system # directories here leaves exactly the libraries that ship with the build. - DD_CORE_DIR="" while read -r _soname _arrow dep _addr; do case "$dep" in /lib/*|/lib64/*|/usr/lib/*|/usr/lib64/*|"") continue ;; esac [[ -f "$dep" ]] || continue - if [[ "$(basename "$dep")" == libdyn_dispatch_core.so* ]]; then - DD_CORE_DIR="$(dirname "$dep")" - fi echo "[home_install] + $(basename "$dep")" install -m 0755 "$dep" "$RAI_LIB_DEST/" done < <(ldd "$CORELIB" | grep '=>' || true) - # libdyn_bins.so is invisible to the ldd walk above: nothing declares it as - # a DT_NEEDED. DynamicDispatch dlopen()s it at runtime, from the directory - # holding libdyn_dispatch_core.so (it dladdr()s itself to find that dir), - # and when the file is absent Transaction::load_large_txn_ops_dll() just - # returns -- no warning, no error. - # - # It only exists when DD was configured with DD_MDS_IN_BINS_DLL=ON, which - # moves every transaction/mds/*.elf -- i.e. the whole AIE_NEXT kernel set -- out - # of the static transaction lib and into this library. Skipping it leaves - # the AIE_NEXT shape table empty, and the first matmul dies with - # Target Shape (K: 3072, N: 3072, Gs: 64) ... not supported in this - # supported shape list. Shape list size: 0 - # which reads like an unsupported model rather than a missing file. - if [[ -n "$DD_CORE_DIR" && -f "$DD_CORE_DIR/libdyn_bins.so" ]]; then - echo "[home_install] + libdyn_bins.so (dlopen'd AIE_NEXT kernel package)" - install -m 0755 "$DD_CORE_DIR/libdyn_bins.so" "$RAI_LIB_DEST/" - fi + # Each model is its own library carrying its kernels, which corelib loads + # only from the directory it was itself loaded from. + for plugin in "$(dirname "$CORELIB")"/*ryzenai_corelib_model_*.so; do + [[ -f "$plugin" ]] || continue + echo "[home_install] + $(basename "$plugin")" + install -m 0755 "$plugin" "$RAI_LIB_DEST/" + done else echo "[home_install] ERROR: --rai, but $BUILD_DIR was not configured with" >&2 echo " FLM_ENABLE_RAI=ON, so there is no corelib runtime to" >&2 @@ -301,8 +286,7 @@ fi # own XRT in the layout XRT expects and reaches it through the binary's RUNPATH, # and every extra hint that was added for rai -- the XILINX_XRT pin, $FLM_PREFIX # /lib ahead of it on LD_LIBRARY_PATH -- only got in its way. A rai prefix has -# the opposite need: no XRT of its own, and corelib plus DynamicDispatch staged -# beside the binary. +# the opposite need: no XRT of its own, and corelib staged beside the binary. if [[ "$WANT_RAI" -eq 1 ]]; then XRT_ENV_BLOCK="$(cat < package, std::shared_ptr runtime, const LM_Config& config, @@ -64,9 +66,11 @@ class gemma4_rai final : public causal_lm { buffer forward(int id) override; /// \brief one multi-token pass: `ids` in, softcapped logits for the last /// position out - /// \note Same path as `forward`. Attention ships no chunked prefill, so a - /// multi-row request at a non-zero position re-runs the tokens behind - /// it together with `ids` from position 0. + /// \note Same path as `forward`. A multi-row pass must start at position 0 + /// or, on a paged cache, on a page boundary, so a multi-row request + /// at any other position re-runs the tokens behind it from the + /// nearest such start together with `ids`: from position 0 unpaged, + /// from the current page's first row paged. /// \throws std::invalid_argument for an empty request, or a continuation /// whose earlier tokens this engine does not hold /// \throws std::out_of_range as `forward` diff --git a/src/include/models/gemma4/rai/aie_next/gemma4_rai_backend.hpp b/src/include/models/gemma4/rai/aie_next/gemma4_rai_backend.hpp index 08e610d7e..31cd48167 100644 --- a/src/include/models/gemma4/rai/aie_next/gemma4_rai_backend.hpp +++ b/src/include/models/gemma4/rai/aie_next/gemma4_rai_backend.hpp @@ -19,25 +19,17 @@ namespace flm::gemma4 { /// \brief the largest context this backend can hold -/// \note THE ENGINE'S OWN CONSTANT, not a second copy of it. `gemma4_rai`'s -/// Impl refuses a `max_length` above `kMaxSequenceLength` and allocates -/// every KV cache and rotary table at exactly that extent, so a traits -/// ceiling written independently could drift above what the engine will -/// accept -- and the frontend checks the TRAITS, before the factory -/// runs, so the drift would surface as a load that passes the frontend's -/// range check and then throws out of the engine constructor. +/// \note THE ENGINE'S OWN CONSTANT. A context in 1..4096 stays on flat_mha. +/// Above that the engine pages at kPageSize, up to kMaxPagedContext. inline constexpr std::uint32_t kRaiContextLimit = - static_cast(kMaxSequenceLength); + static_cast(kMaxPagedContext); -/// \brief the largest number of tokens corelib will decode -/// \note Likewise `kMaxDecodeWindow`, one below the context: the last position -/// in the window has to be readable, so a decode that filled it would -/// have nowhere to put the next token. +/// \brief one below the longest context: the last position has to be readable inline constexpr std::uint32_t kRaiDecodeLimit = - static_cast(kMaxDecodeWindow); + static_cast(kMaxPagedContext - 1); /// \brief what the frontend must know before building this backend -/// \return no xclbin, no preemption, 4096-token context ceiling +/// \return no xclbin, no preemption, 16384-token context ceiling /// \note Pure data, and inline on purpose: `AutoModel::_shared_load_backend` /// reads it while assembling the BackendContext -- that is, BEFORE any /// engine exists -- so it cannot be a virtual on ModelBackend. Being in diff --git a/src/include/models/gemma4/rai/aie_next/gemma4_rai_constants.hpp b/src/include/models/gemma4/rai/aie_next/gemma4_rai_constants.hpp index ffc7cad41..d1c324b21 100644 --- a/src/include/models/gemma4/rai/aie_next/gemma4_rai_constants.hpp +++ b/src/include/models/gemma4/rai/aie_next/gemma4_rai_constants.hpp @@ -15,6 +15,9 @@ namespace flm::gemma4 { /// reasoning applied to Phi-4's own (different) pair. inline constexpr int kPrefillPdi = 8; inline constexpr int kTokenPdi = 17; +/// Decode PDI for a context longer than one attention window. Prefill stays +/// on kPrefillPdi; the paged kernels ship on this overlay and nowhere else. +inline constexpr int kPagedTokenPdi = 33; /// Shared by both shipped rows (E2B and E4B) -- read from the GGUF's own /// `tokenizer.ggml.tokens` array length in a real load, but the two models @@ -27,8 +30,26 @@ inline constexpr std::int64_t kVocabularySize = 262144; /// group fields below it were established. inline constexpr std::uint32_t kGroupSize = 32; +/// Longest context flat_mha ships, and the page a longer context is cut into. +inline constexpr std::int64_t kPageSize = 4096; inline constexpr std::int64_t kMaxSequenceLength = 4096; inline constexpr std::int64_t kMaxDecodeWindow = 4095; +/// Longest paged context: 16384, the validated ceiling. corelib upstream +/// validated paged Gemma 4 to 16k (aie4-models.md, "Past 4096 tokens"), well +/// short of the model's own 131072. A GGUF stating a shorter +/// `gemma4.context_length` caps it further. Whole pages. +inline constexpr std::int64_t kMaxPagedContext = 16384; +/// Retries for a non-finite paged prefill chunk past page 0. The chunk +/// rewrites only its own page, so a repeat is the same work. +inline constexpr int kPagedChunkAttempts = 3; + +/// \return true when `length` is 1..4096, or a multiple of kPageSize up to +/// kMaxPagedContext +inline bool ContextFits(std::uint32_t length) { + if (length < 1 || length > static_cast(kMaxPagedContext)) return false; + if (length <= static_cast(kMaxSequenceLength)) return true; + return length % static_cast(kPageSize) == 0; +} /// Applied to the final logits on every Gemma 4 size: tanh(x / cap) * cap. inline constexpr float kLogitSoftcap = 30.0f; @@ -41,7 +62,7 @@ inline constexpr float kLogitSoftcap = 30.0f; inline constexpr float kRmsEpsilon = 1.0e-6f; /// See phi4_rai_constants.hpp: 0 lets corelib choose the per-create packer -/// threads, on top of the concurrent creates below. +/// threads (8 at 0.25.0), on top of the concurrent creates below. inline constexpr std::uint32_t kRequantizeThreads = 0; inline constexpr std::size_t kWeightCreateConcurrency = 8; diff --git a/src/include/models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp b/src/include/models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp index 544ef5190..5a1ef0f45 100644 --- a/src/include/models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp +++ b/src/include/models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp @@ -1,5 +1,7 @@ #pragma once +#include "models/gemma4/rai/aie_next/gemma4_rai_constants.hpp" + #include #include #include @@ -194,6 +196,7 @@ struct Gemma4Config { std::int32_t bos_token_id; bool add_bos; std::uint32_t group, head_group, ple_group; // from the row + std::int64_t context_length = kMaxPagedContext; // gemma4.context_length /// \brief the head size a layer's attention runs at -- `head_dim` if it /// slides, `global_head_dim` if it is full-attention diff --git a/src/include/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.hpp b/src/include/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.hpp index f145425ac..eefdc72af 100644 --- a/src/include/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.hpp +++ b/src/include/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.hpp @@ -1,5 +1,6 @@ #pragma once +#include "models/gemma4/rai/aie_next/gemma4_rai_constants.hpp" #include "models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp" #include "rai/corelib_api.hpp" @@ -77,9 +78,14 @@ class Gemma4ShapePlan final { /// pair a stream was opened with selects the ELF set, and a shape /// exists under one pair and not another /// \param config the model's shape, read off its GGUF (Task C5) + /// \param context the engine's context. Above kPageSize the plan is paged: + /// attention descriptors carry that window and a kPageSize page. static Gemma4ShapePlan Build( const std::shared_ptr& api, - ryzenai_corelib_stream_ptr stream, const Gemma4Config& config); + ryzenai_corelib_stream_ptr stream, const Gemma4Config& config, + std::int64_t context = kMaxSequenceLength); + + bool paged() const noexcept { return paged_; } /// \param live_rows the live (unpadded) row count a caller wants extents /// for, 1..4096 @@ -105,6 +111,11 @@ class Gemma4ShapePlan final { const ryzenai_corelib_flat_mha_bf16_desc& attention_desc( bool swa, bool shared) const noexcept; + /// \note Valid only when paged() is true. Same (geometry, cache role) + /// index as attention_desc. + const ryzenai_corelib_paged_mha_bf16_desc& paged_attention_desc( + bool swa, bool shared) const noexcept; + /// \brief the padded row extent lm_head's matmul runs at -- the leading /// extent of the `[rows, vocab]` logits tensor a caller allocates /// \note ONE, on every shipped Gemma 4, and NOT a simplification. lm_head @@ -135,6 +146,8 @@ class Gemma4ShapePlan final { /// \note Materialized once in Build() and never written again. See the /// class comment for why four and not two. std::array attention_descs_{}; + std::array paged_descs_{}; + bool paged_{false}; }; } // namespace flm::gemma4 diff --git a/src/include/models/phi4/rai/aie_next/phi4_rai_constants.hpp b/src/include/models/phi4/rai/aie_next/phi4_rai_constants.hpp index fe575972f..430db8aa2 100644 --- a/src/include/models/phi4/rai/aie_next/phi4_rai_constants.hpp +++ b/src/include/models/phi4/rai/aie_next/phi4_rai_constants.hpp @@ -19,7 +19,7 @@ inline constexpr std::int64_t kMaxDecodeWindow = 4095; inline constexpr std::uint32_t kRequantizedGroupSize = 64; /// Intra-packer thread hint for one Q8_0 requantizing create, passed as /// weights_options.fast_packer_threads. 0 lets corelib choose (8 threads at -/// 0.11.0), so with kWeightCreateConcurrency creates in flight the packer +/// 0.25.0), so with kWeightCreateConcurrency creates in flight the packer /// threads multiply; 1 keeps each create serial. inline constexpr std::uint32_t kRequantizeThreads = 0; diff --git a/src/include/models/qwen36/rai/aie_next/qwen36_rai.hpp b/src/include/models/qwen36/rai/aie_next/qwen36_rai.hpp new file mode 100644 index 000000000..585f58d5c --- /dev/null +++ b/src/include/models/qwen36/rai/aie_next/qwen36_rai.hpp @@ -0,0 +1,46 @@ +/// \file qwen36_rai.hpp +/// \brief Qwen3.6-35B-A3B text on ryzenai-corelib, from a Q8_0 GGUF +#pragma once + +#include "causal_lm.hpp" +#include "rai/corelib_runtime.hpp" + +#include +#include +#include + +namespace flm::qwen36 { + +inline constexpr std::uint32_t kContextLimit = 262144; + +/// \brief Qwen3.6 text through corelib's model API +/// \note Vision is not this engine. A prompt is one forward(); corelib runs +/// its first 4096 rows from position 0 as one pass and steps the rest. +/// The KV cache is sized once, at the capacity `max_length` rounds up +/// to in whole 4096-row pages. +class qwen36_rai final : public causal_lm { +public: + qwen36_rai(std::filesystem::path gguf, + std::shared_ptr runtime, + std::uint32_t max_length = 4096); + ~qwen36_rai() override; + + buffer forward(int id) override; + buffer prefill(std::vector& ids, void* payload = nullptr) override; + void set_context_length(int length) override; + void load_weights(Q4NX&) override; + void update_max_length(std::uint32_t max_length) override; + void clear_context() override; + buffer get_k_cache(int layer, int index) override; + buffer get_v_cache(int layer, int index) override; + int get_current_context_length() override; + int checkpoint() override; + int restore() override; + bool poisoned() const noexcept; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace flm::qwen36 diff --git a/src/include/models/qwen36/rai/aie_next/qwen36_rai_backend.hpp b/src/include/models/qwen36/rai/aie_next/qwen36_rai_backend.hpp new file mode 100644 index 000000000..5665ae526 --- /dev/null +++ b/src/include/models/qwen36/rai/aie_next/qwen36_rai_backend.hpp @@ -0,0 +1,25 @@ +/// \file qwen36_rai_backend.hpp +/// \brief The ryzenai-corelib backend for Qwen3.6-35B-A3B text +/// \note Compiled only when FLM_ENABLE_RAI is on. +#pragma once + +#include "AutoModel/model_backend.hpp" + +#include + +namespace flm::qwen36 { + +inline constexpr std::uint32_t kRaiContextLimit = 262144; +inline constexpr std::uint32_t kRaiDecodeLimit = kRaiContextLimit - 1; + +inline flm::backend::BackendTraits rai_traits() { + flm::backend::BackendTraits traits; + traits.needs_npu_xclbin = false; + traits.supports_preemption = false; + traits.max_context_length = kRaiContextLimit; + return traits; +} + +flm::backend::BackendFactory rai_factory(); + +} // namespace flm::qwen36 diff --git a/src/include/rai/corelib_api.hpp b/src/include/rai/corelib_api.hpp index 7816a7a25..2c57f4ab9 100644 --- a/src/include/rai/corelib_api.hpp +++ b/src/include/rai/corelib_api.hpp @@ -2,9 +2,9 @@ #include -#if RYZENAI_CORELIB_VERSION_MAJOR != 0 || RYZENAI_CORELIB_VERSION_MINOR != 11 || \ +#if RYZENAI_CORELIB_VERSION_MAJOR != 0 || RYZENAI_CORELIB_VERSION_MINOR != 26 || \ RYZENAI_CORELIB_VERSION_PATCH != 0 -#error "FastFlowLM requires ryzenai-corelib headers exactly 0.11.0" +#error "FastFlowLM requires ryzenai-corelib headers exactly 0.26.0" #endif #include @@ -35,6 +35,7 @@ X(tensor_get_data_type, ryzenai_corelib_tensor_get_data_type) \ X(stream_get_kernels_root, ryzenai_corelib_stream_get_kernels_root) \ X(stream_set_kernels_root, ryzenai_corelib_stream_set_kernels_root) \ + X(stream_enable_kernel_bundle, ryzenai_corelib_stream_enable_kernel_bundle) \ X(weights_slice, ryzenai_corelib_weights_slice) \ X(matmul_enum_kernels, ryzenai_corelib_matmul_bf16_enum_kernels) \ X(matmul_weights_create, ryzenai_corelib_matmul_bf16_weights_create) \ @@ -51,9 +52,18 @@ X(ple, ryzenai_corelib_ple_bf16) \ X(flat_mha_enum_kernels, ryzenai_corelib_flat_mha_bf16_enum_kernels) \ X(flat_mha, ryzenai_corelib_flat_mha_bf16) \ + X(paged_mha_enum_kernels, ryzenai_corelib_paged_mha_bf16_enum_kernels) \ + X(paged_mha, ryzenai_corelib_paged_mha_bf16) \ X(linear_attention, ryzenai_corelib_linear_attention_bf16) \ X(dwconv_weights_create, ryzenai_corelib_dwconv_bf16_weights_create) \ X(dwconv, ryzenai_corelib_dwconv_bf16) \ + X(model_info_get, ryzenai_corelib_model_info_get) \ + X(model_forward, ryzenai_corelib_model_forward) \ + X(model_reset, ryzenai_corelib_model_reset) \ + X(model_weights_create_path, ryzenai_corelib_model_weights_create_path) \ + X(model_create, ryzenai_corelib_model_create) \ + X(model_state_save, ryzenai_corelib_model_state_save) \ + X(model_state_restore, ryzenai_corelib_model_state_restore) \ X(cleanup, ryzenai_corelib_cleanup) namespace flm::corelib { @@ -96,10 +106,17 @@ class CorelibApi final { #if defined(FLM_CORELIB_LINK_STATIC) /// \brief bind to the corelib linked into this binary /// \return the API bound to the linked symbols - /// \note The static build has no DLL to pick, so FLM_RAI_CORELIB_PATH has - /// nothing to select and is reported as ignored. + /// \note The static build has no library to pick, so FLM_RAI_CORELIB_PATH + /// has nothing to select and is reported as ignored. On Windows an + /// import library still leaves a DLL that the loader binds before + /// main, so loaded_library_path() reports the module the linked + /// symbols live in. static std::shared_ptr LoadStatic(); #endif + /// \brief the file of the module this process has loaded under `name` + /// \return an absolute path, or empty when no such module is loaded or the + /// platform has no way to ask + static std::filesystem::path LoadedModulePath(const wchar_t* name); static std::shared_ptr ResolveForTest( Resolver resolver, std::filesystem::path loaded_library_path = {}); static std::filesystem::path ResolveLibraryPath( diff --git a/src/include/rai/corelib_object.hpp b/src/include/rai/corelib_object.hpp index 2558a99ce..5f7a72da3 100644 --- a/src/include/rai/corelib_object.hpp +++ b/src/include/rai/corelib_object.hpp @@ -3,6 +3,7 @@ #include "rai/corelib_api.hpp" #include +#include #include namespace flm::corelib { @@ -16,6 +17,9 @@ struct SsMlpWeightsTag {}; struct RmsNormWeightsTag {}; struct PleWeightsTag {}; struct DwConvWeightsTag {}; +struct ModelTag {}; +struct ModelWeightsTag {}; +struct ModelSnapshotTag {}; template class UniqueObject final { @@ -68,5 +72,23 @@ using UniqueSsMlpWeights = UniqueObject; using UniqueRmsNormWeights = UniqueObject; using UniquePleWeights = UniqueObject; using UniqueDwConvWeights = UniqueObject; +using UniqueModel = UniqueObject; +using UniqueModelWeights = UniqueObject; +using UniqueModelSnapshot = UniqueObject; + +/// \brief a stream on the PDI pair that resolves the kernels `bundle` ships +/// \note A new stream resolves no kernel at all: every enumeration comes back +/// empty and every dispatch runs on the host until a bundle is enabled, +/// so this has to precede the shape plan and the weights. +inline UniqueStream OpenModelStream(const std::shared_ptr& api, + int prefill_pdi, int token_pdi, const char* bundle) { + void* raw = nullptr; + api->Check(api->functions().create_stream(prefill_pdi, token_pdi, &raw), + "ryzenai_corelib_create_stream"); + UniqueStream stream(api, raw); + api->Check(api->functions().stream_enable_kernel_bundle(stream.get(), bundle), + std::string("ryzenai_corelib_stream_enable_kernel_bundle ") + bundle); + return stream; +} } // namespace flm::corelib diff --git a/src/include/rai/host_ops.hpp b/src/include/rai/host_ops.hpp index 73c72e891..92d2de560 100644 --- a/src/include/rai/host_ops.hpp +++ b/src/include/rai/host_ops.hpp @@ -23,4 +23,7 @@ void RmsNorm(std::span input, std::span scale, /// \brief round-to-nearest-even F32 -> BF16, NaN kept quiet std::vector F32ToBf16(std::span values); +/// \brief the same conversion into a caller's buffer of values.size() elements +void F32ToBf16(std::span values, std::span output); + } // namespace flm::rai diff --git a/src/include/rai/kernel_grid.hpp b/src/include/rai/kernel_grid.hpp index 1c046f306..9d0341194 100644 --- a/src/include/rai/kernel_grid.hpp +++ b/src/include/rai/kernel_grid.hpp @@ -1,6 +1,6 @@ /// \file kernel_grid.hpp /// \brief The shapes a corelib stream can actually run -/// \note corelib 0.11 reports shipped kernels and rounds nothing. A reported +/// \note corelib reports shipped kernels and rounds nothing. A reported /// M below zero is a test double meaning "every requested M ships /// exactly"; a real corelib never says that. #pragma once @@ -99,6 +99,31 @@ class ShapeGrid { FlatCtx ctx{&grid, &desc}; api.Check(api.functions().flat_mha_enum_kernels(stream, &OnFlatMha, &ctx), "ryzenai_corelib_flat_mha_bf16_enum_kernels"); + // A decode step is seq_len_q == 1, and that ELF has six numeric fields. + // enum_kernels declines the ambiguous ones, so the callback can miss + // every token kernel and the smallest reported prefill M is 64. + // flat_mha still dispatches M=1 and puts the KV bucket in seq_pad + // (the 64 in ..._8_1_1_64_256_4096_scale1). Record the row, or a + // one-token step is rounded up to that prefill kernel and refused. + if (!grid.wildcard && !grid.Has(1, 0, 0, -1)) grid.shapes.push_back({1, 0, 0, 0}); + return grid; + } + + /// \brief paged attention kernels whose page matches `desc.page_size` + static ShapeGrid PagedMha(const CorelibApi& api, ryzenai_corelib_stream_ptr stream, + const ryzenai_corelib_paged_mha_bf16_desc& desc) { + ShapeGrid grid; + PagedCtx ctx{&grid, &desc}; + api.Check(api.functions().paged_mha_enum_kernels(stream, &OnPagedMha, &ctx), + "ryzenai_corelib_paged_mha_bf16_enum_kernels"); + // A decode step is seq_len_q == 1, and paged_mha puts the KV bucket in + // seq_pad: ..._8_1_1_64_256_4096, where 1 is the query rows and 64 is + // seq_pad. corelib 0.25 reads that six-field name as a token kernel + // and walks the decode PDI, so Gemma 4's p33 decode kernel is + // reported. A corelib that does not report M=1 would round a decode + // up to the smallest prefill kernel (64) and refuse it, so M=1 is + // still recorded; Has() keeps it from being added twice. + if (!grid.wildcard && !grid.Has(1, 0, 0, -1)) grid.shapes.push_back({1, 0, 0, 0}); return grid; } @@ -111,6 +136,10 @@ class ShapeGrid { ShapeGrid* grid; const ryzenai_corelib_flat_mha_bf16_desc* desc; }; + struct PagedCtx { + ShapeGrid* grid; + const ryzenai_corelib_paged_mha_bf16_desc* desc; + }; static bool OnMatmul(void* ctx, int, std::int64_t m, std::int64_t k, std::int64_t n, std::int64_t group, bool, std::int64_t, std::int64_t) { @@ -153,6 +182,20 @@ class ShapeGrid { pack->grid->shapes.push_back({seq_len_q, 0, 0, 0}); return true; } + + static bool OnPagedMha(void* ctx, int, std::int64_t heads, std::int64_t kv_heads, + std::int64_t seq_len_q, std::int64_t, std::int64_t head_size, + std::int64_t page_size, std::int64_t) { + auto* pack = static_cast(ctx); + if (seq_len_q < 0) pack->grid->wildcard = true; + const auto* desc = pack->desc; + if (seq_len_q >= 0 && + (heads != desc->num_heads || kv_heads != desc->kv_num_heads || + head_size != desc->head_size || page_size != desc->page_size)) + return true; + pack->grid->shapes.push_back({seq_len_q, 0, 0, 0}); + return true; + } }; /// \brief the row count to allocate and dispatch for one requested bucket diff --git a/src/model_info.json b/src/model_info.json index 8981d61b9..742418395 100644 --- a/src/model_info.json +++ b/src/model_info.json @@ -4070,5 +4070,37 @@ "size": 3126, "sha256": "d0883072e01861ed0b2d47be3c16c36a8e81c224c7ffaa310c6558fb3f932b05" } + ], + "qwen3.6-moe-rai:35b-a3b": [ + { + "type": "file", + "path": "Qwen3.6-35B-A3B-Q8_0.gguf", + "size": 36903140320, + "sha256": "d1a395809f65a43a13ad119eb4e7acdef1ac6d68120f39902c8ab96e72794a59" + }, + { + "type": "file", + "path": "tokenizer.json", + "size": 12807982, + "sha256": "5f9e4d4901a92b997e463c1f46055088b6cca5ca61a6522d1b9f64c4bb81cb42" + }, + { + "type": "file", + "path": "tokenizer_config.json", + "size": 16718, + "sha256": "5186f0defcd7f232382c7f0aebcd2252d073bb921ab240e407b7ae8745d2b29b" + }, + { + "type": "file", + "path": "config.json", + "size": 3686, + "sha256": "93a4693fa9d8392fbfccd4b3c9873f4bfdcb14fdede978b123d07d19675efe99" + }, + { + "type": "file", + "path": "chat_template.jinja", + "size": 7764, + "sha256": "e84f32a23fdda27689f868aa4a1a5621f41133e51a48d7f3efcbea2839574259" + } ] } diff --git a/src/model_list.json b/src/model_list.json index 3aceb02a4..5c77ac374 100644 --- a/src/model_list.json +++ b/src/model_list.json @@ -1451,7 +1451,7 @@ "url": "https://huggingface.co/unsloth/gemma-4-E2B-it-GGUF/resolve/0314792d7f1f7e229411f620751375812bb9faf2", "file_url": "https://huggingface.co/api/models/unsloth/gemma-4-E2B-it-GGUF/tree/0314792d7f1f7e229411f620751375812bb9faf2", "size": 5080530526, - "default_context_length": 4096, + "default_context_length": 16384, "max_prefill_len": 4096, "details": { "family": "gemma4e", @@ -1502,7 +1502,7 @@ "url": "https://huggingface.co/unsloth/gemma-4-E4B-it-GGUF/resolve/bfc15c382204943c3a8fff0c750b94ae2364d7a3", "file_url": "https://huggingface.co/api/models/unsloth/gemma-4-E4B-it-GGUF/tree/bfc15c382204943c3a8fff0c750b94ae2364d7a3", "size": 8225131325, - "default_context_length": 4096, + "default_context_length": 16384, "max_prefill_len": 4096, "details": { "family": "gemma4e", @@ -1592,6 +1592,58 @@ "tool-calling" ] } + }, + "qwen3.6-moe-rai": { + "35b-a3b": { + "supported_platforms": [ + "aie_next" + ], + "name": "Qwen3.6-35B-A3B", + "url": "https://huggingface.co/unsloth/Qwen3.6-35B-A3B-GGUF/resolve/a483e9e6cbd595906af30beda3187c2663a1118c", + "file_url": "https://huggingface.co/api/models/unsloth/Qwen3.6-35B-A3B-GGUF/tree/a483e9e6cbd595906af30beda3187c2663a1118c", + "size": 36915976470, + "default_context_length": 4096, + "max_prefill_len": 4096, + "details": { + "family": "qwen3.6-moe", + "think": true, + "parameter_size": "35B", + "quantization_level": "Q8_0 -> group-32" + }, + "flm_min_version": "1.0.6", + "vlm": false, + "files": [ + "Qwen3.6-35B-A3B-Q8_0.gguf", + "tokenizer.json", + "tokenizer_config.json", + "config.json", + "chat_template.jinja" + ], + "file_sources": { + "tokenizer.json": { + "url": "https://huggingface.co/Qwen/Qwen3.6-35B-A3B", + "revision": "995ad96eacd98c81ed38be0c5b274b04031597b0" + }, + "tokenizer_config.json": { + "url": "https://huggingface.co/Qwen/Qwen3.6-35B-A3B", + "revision": "995ad96eacd98c81ed38be0c5b274b04031597b0" + }, + "config.json": { + "url": "https://huggingface.co/Qwen/Qwen3.6-35B-A3B", + "revision": "995ad96eacd98c81ed38be0c5b274b04031597b0" + }, + "chat_template.jinja": { + "url": "https://huggingface.co/Qwen/Qwen3.6-35B-A3B", + "revision": "995ad96eacd98c81ed38be0c5b274b04031597b0" + } + }, + "footprint": 36.9, + "model_info_key": "qwen3.6-moe-rai:35b-a3b", + "label": [ + "reasoning", + "tool-calling" + ] + } }, "qwen3-rai": { "0.6b": { diff --git a/src/test/gemma4_rai/fake_corelib.cpp b/src/test/gemma4_rai/fake_corelib.cpp index 2b1f8fb5d..e9a3e5613 100644 --- a/src/test/gemma4_rai/fake_corelib.cpp +++ b/src/test/gemma4_rai/fake_corelib.cpp @@ -31,6 +31,10 @@ fake_corelib::State state; /// nothing here promises one entry point never calls another. std::recursive_mutex state_mutex; +/// \brief what `get_last_error_message` reports; set only by refusals that +/// carry corelib's own wording +std::string last_error; + /// \brief the backing allocation a tensor and all its windows share /// /// Held by shared_ptr so that a window (and a window of a window) keeps the @@ -48,6 +52,8 @@ struct FakeObject { /// \brief offset from the ROOT allocation, in elements, accumulated /// through however many levels of window std::size_t window_offset{}; + /// \brief the strides a window was created with, empty for dense + std::vector strides; std::shared_ptr storage; /// \brief the packed N a matmul weights object was created for /// \note What lets a matmul dispatch pick the SAME `"matmul-"` helper @@ -400,7 +406,7 @@ struct TypedFake { } else if constexpr (std::is_same_v) { return "success"; } else if constexpr (std::is_same_v) { - return ""; + return last_error.c_str(); } else if constexpr (std::is_same_v) { static const char kKernelsRoot[] = ""; return kKernelsRoot; @@ -462,12 +468,17 @@ struct TypedFake { *out = object; state.host_view_creates.push_back({type, object->shape, object}); return ryzenai_corelib_status_success; + } else if constexpr (std::is_same_v) { + const char* model = std::get<1>(arguments); + state.kernel_bundles.push_back( + {std::get<0>(arguments), model ? model : "", state.enum_calls.size()}); + return ryzenai_corelib_status_success; } else if constexpr (std::is_same_v) { void* parent = std::get<0>(arguments); const auto* shape = std::get<1>(arguments); const auto shape_len = std::get<2>(arguments); const auto offset = std::get<3>(arguments); - auto* out = std::get<4>(arguments); + auto* out = std::get<6>(arguments); if (out) *out = nullptr; if (!out || !shape) return ryzenai_corelib_status_bad_argument; ryzenai_corelib_data_type data_type = ryzenai_corelib_data_type_bf16; @@ -498,8 +509,12 @@ struct TypedFake { object->window_offset = window_offset; object->shape = window_shape; object->byte_size = Elements(object->shape) * TypeBytes(object->data_type); + const auto* strides = std::get<4>(arguments); + const auto strides_len = std::get<5>(arguments); + if (strides) object->strides.assign(strides, strides + strides_len); *out = object; - state.tensor_windows.push_back({parent, object->shape, offset, object}); + state.tensor_windows.push_back( + {parent, object->shape, offset, object, object->strides}); return ryzenai_corelib_status_success; } else if constexpr (std::is_same_v) { auto* object = static_cast(std::get<0>(arguments)); @@ -580,6 +595,20 @@ struct TypedFake { } } return ryzenai_corelib_status_success; + } else if constexpr (std::is_same_v) { + // A context above one page plans against this grid, not + // flat_mha's, so an empty answer leaves every prefill bucket + // without a kernel. + auto fn = std::get<1>(arguments); + if (fn && state.shipped_paged.empty()) { + fn(std::get<2>(arguments), 0, 0, 0, -1, 0, 0, 0, 0); + } else if (fn) { + for (const auto& shipped : state.shipped_paged) { + fn(std::get<2>(arguments), 0, shipped.heads, shipped.kv_heads, shipped.m, 0, + shipped.head_size, shipped.page_size, 0); + } + } + return ryzenai_corelib_status_success; } else if constexpr (std::is_same_v) { // 0.5.0: (stream, m*, k*, n*, group_size). Every padding helper // gained a leading stream in 0.5.0 -- recorded here so @@ -910,7 +939,8 @@ struct TypedFake { std::is_same_v || std::is_same_v || std::is_same_v || - std::is_same_v) { + std::is_same_v || + std::is_same_v) { // THE INJECTED MID-CHAIN REFUSAL, BEFORE ANY VALIDATION AND // BEFORE THE RECORD. See State::fail_dispatch_at: this is the // only way a test can reach the engine's catch block, which is @@ -957,6 +987,20 @@ struct TypedFake { if (!RowsAgree({input, output}) || !OnGrid(helper, input) || !OnGrid(helper, output)) return ryzenai_corelib_status_bad_argument; + // Like corelib: a 3-D output is the cache-writing form and + // must be a strided window [k * head, head, 1] with k >= its + // rows. A dense window over the whole cache is refused too. + if (output && RankOf(output) == 3) { + const auto* object = static_cast(output); + const auto& strides = object->strides; + const auto head = object->shape[2]; + if (head <= 0 || strides.size() != 3 || strides[2] != 1 || strides[1] != head || + strides[0] % head != 0 || strides[0] / head < object->shape[1]) { + last_error = "a cache-writing 3-D output must be a window with " + "strides [max_seq * head_size, head_size, 1]"; + return ryzenai_corelib_status_bad_argument; + } + } } else if constexpr (std::is_same_v) { void* input = std::get<1>(arguments); void* residual = std::get<2>(arguments); @@ -1037,8 +1081,10 @@ struct TypedFake { // one -- but recording it is what makes the reference // driver's load-time guard (gemma4_driver.py:1670-1690) // portable to a test. - record.has_mha_desc = desc != nullptr; - if (desc) record.mha_desc = *desc; + if constexpr (std::is_same_v) { + record.has_mha_desc = desc != nullptr; + if (desc) record.mha_desc = *desc; + } // THE ROTARY PAIR AND THE TWO CACHES. Not row-bearing, so not // in `operands` and not checked here -- retained so that "this // layer roped with the sliding table and attended over cache @@ -1054,15 +1100,21 @@ struct TypedFake { if (!table || static_cast(table)->kind != "host_view") return ryzenai_corelib_status_bad_argument; } - if (!RowsAgree({q, k, out}) || !OnGrid("mha", q) || - !OnGrid("mha", k) || !OnGrid("mha", out)) + // paged_mha has no pad helper, so there is no answer set to + // put its rows on. + if (!RowsAgree({q, k, out})) return ryzenai_corelib_status_bad_argument; + if constexpr (std::is_same_v) { + if (!OnGrid("mha", q) || !OnGrid("mha", k) || !OnGrid("mha", out)) + return ryzenai_corelib_status_bad_argument; + } } - record.kind = std::is_same_v ? "matmul" - : std::is_same_v ? "ssmlp" - : std::is_same_v ? "rmsnorm" - : std::is_same_v ? "ple" - : "mha"; + record.kind = std::is_same_v ? "matmul" + : std::is_same_v ? "ssmlp" + : std::is_same_v ? "rmsnorm" + : std::is_same_v ? "ple" + : std::is_same_v ? "paged_mha" + : "mha"; record.stream = std::get<0>(arguments); record.operands = operands; record.operand_storage.reserve(operands.size()); @@ -1158,10 +1210,12 @@ ryzenai_corelib_status NoteMhaRows(ryzenai_corelib_stream_ptr stream, void Reset() { std::lock_guard lock(state_mutex); state.enum_calls.clear(); + state.kernel_bundles.clear(); state.explicit_grid = false; state.shipped_matmul.clear(); state.shipped_ssmlp.clear(); state.shipped_mha.clear(); + state.shipped_paged.clear(); state.matmul_pad_calls.clear(); state.rows_pad_calls.clear(); state.mha_pad_calls.clear(); @@ -1184,6 +1238,7 @@ void Reset() { state.fail_dispatch_at.reset(); state.read_returns_nonfinite = false; state.nonfinite_read_index.reset(); + last_error.clear(); } flm::corelib::CorelibApi::Resolver Resolver() { diff --git a/src/test/gemma4_rai/fake_corelib.hpp b/src/test/gemma4_rai/fake_corelib.hpp index a506060e5..eb36b6d58 100644 --- a/src/test/gemma4_rai/fake_corelib.hpp +++ b/src/test/gemma4_rai/fake_corelib.hpp @@ -306,6 +306,11 @@ struct ShippedMha { bool kv_shared, scale_one; }; +/// \brief one paged attention kernel `paged_mha_enum_kernels` reports +struct ShippedPaged { + std::int64_t heads, kv_heads, m, head_size, page_size; +}; + /// \brief one `ryzenai_corelib_matmul_bf16_weights_create_*` call /// \note The WHOLE descriptor, because `group_size` is the field C7 must get /// right and it is the one nothing else in this codebase can observe: @@ -490,6 +495,8 @@ struct TensorWindowRecord { std::vector shape; std::size_t offset; void* object; + /// \brief the strides the window was created with, empty for dense + std::vector strides; }; /// \brief one `ryzenai_corelib_tensor_write` call @@ -685,6 +692,13 @@ struct DispatchRecord { void* weights{nullptr}; }; +/// \brief one `ryzenai_corelib_stream_enable_kernel_bundle` call +struct KernelBundleRecord { + void* stream{}; + std::string model; + std::size_t enum_calls_before{}; +}; + /// \note EVERY MEMBER IS WRITTEN UNDER ONE MUTEX inside the fake, because /// the engine packs its weights across a pool of threads /// (`kWeightCreateConcurrency`). A reader -- i.e. a test -- runs after @@ -700,6 +714,7 @@ struct DispatchRecord { struct State { /// \brief every `..._enum_kernels` call, in order std::vector enum_calls; + std::vector kernel_bundles; /// \brief report exactly the `shipped_*` kernels instead of "every M" /// \note What makes a plan's kernel QUERY observable: the plan filters /// the reported set itself, so the only thing a test can see is @@ -708,6 +723,8 @@ struct State { std::vector shipped_matmul; std::vector shipped_ssmlp; std::vector shipped_mha; + /// \brief report exactly these paged kernels; empty reports "every M" + std::vector shipped_paged; std::vector matmul_pad_calls; std::vector rows_pad_calls; std::vector mha_pad_calls; diff --git a/src/test/gemma4_rai/gemma4_gguf_fixture.hpp b/src/test/gemma4_rai/gemma4_gguf_fixture.hpp index 6108a2c8b..f9af43f76 100644 --- a/src/test/gemma4_rai/gemma4_gguf_fixture.hpp +++ b/src/test/gemma4_rai/gemma4_gguf_fixture.hpp @@ -734,6 +734,8 @@ inline Builder MakeBuilder(const FixtureOptions& options) { builder.SetMetadata("gemma4.attention.shared_kv_layers", static_cast(options.shared_kv_layers)); builder.SetMetadata("gemma4.attention.sliding_window", std::uint32_t{512}); + // Both shipped files state 131072. + builder.SetMetadata("gemma4.context_length", std::uint32_t{131072}); // 9.999999974752427e-07f is the exact float32 bit pattern the real files // carry for 1e-6 -- see verified-gguf-facts.md's "one value the plan // never stated". Not Phi-4's 1e-5 (phi4_rai_constants.hpp's kRmsEpsilon). diff --git a/src/test/gemma4_rai/test_gemma4_backend.cpp b/src/test/gemma4_rai/test_gemma4_backend.cpp index ebc7f83f3..ee0bd856d 100644 --- a/src/test/gemma4_rai/test_gemma4_backend.cpp +++ b/src/test/gemma4_rai/test_gemma4_backend.cpp @@ -89,6 +89,9 @@ constexpr const char* kRai = "rai"; constexpr int kRealEosId = 106; constexpr const char* kRealEosToken = ""; +/// \brief the other stop id real packages list in config.json (``) +constexpr int kRealConfigEosId = 1; + /// \brief the BOS id every real Gemma 4 package states, from three sources /// \note `tokenizer.ggml.bos_token_id` in the GGUF, `text_config.bos_token_id` /// in config.json -- NOT the top level, which is where the EOS lives and @@ -423,7 +426,7 @@ void TestTraitsAreReadableWithoutConstructingAnEngine() { const auto traits = flm::gemma4::rai_traits(); TEST_REQUIRE(!traits.needs_npu_xclbin); TEST_REQUIRE(!traits.supports_preemption); - TEST_REQUIRE(traits.max_context_length == 4096); + TEST_REQUIRE(traits.max_context_length == 16384); // Nothing above touched corelib, a GGUF or the filesystem, which is the // property that lets the frontend consult it while it is still ASSEMBLING // the BackendContext -- before any engine exists to ask. @@ -439,18 +442,24 @@ void TestTraitsComeBackOutOfTheRegistryUnchanged() { } void TestTheTraitsCeilingIsTheEnginesOwnLimit() { - // Two numbers that must not drift apart: the frontend rejects an - // over-long context against the TRAITS, and the engine rejects it again - // against kMaxSequenceLength. A traits ceiling above the engine's would - // pass the frontend and throw out of the engine constructor instead. + // The frontend rejects an over-long context against the TRAITS. The + // engine accepts the same ceiling: 1..4096 unpaged, or whole pages up to + // kMaxPagedContext. TEST_REQUIRE(flm::gemma4::kRaiContextLimit == - static_cast(flm::gemma4::kMaxSequenceLength)); - TEST_REQUIRE(flm::gemma4::kRaiDecodeLimit == - static_cast(flm::gemma4::kMaxDecodeWindow)); + static_cast(flm::gemma4::kMaxPagedContext)); TEST_REQUIRE(flm::gemma4::kRaiDecodeLimit + 1 == flm::gemma4::kRaiContextLimit); } +void TestContextFitsStopsAtTheValidated16384() { + TEST_REQUIRE(flm::gemma4::ContextFits(4096)); + TEST_REQUIRE(flm::gemma4::ContextFits(8192)); + TEST_REQUIRE(flm::gemma4::ContextFits(16384)); + TEST_REQUIRE(!flm::gemma4::ContextFits(10000)); + TEST_REQUIRE(!flm::gemma4::ContextFits(20480)); + TEST_REQUIRE(!flm::gemma4::ContextFits(131072)); +} + // --------------------------------------------------------------------------- // 2. Everything is refused before anything is opened or allocated. @@ -480,6 +489,17 @@ void TestAnOverLongContextIsRejectedBeforeAnythingIsOpened() { TEST_REQUIRE(runtime.nothing_created()); } +void TestAContextPastThePagedCeilingIsRejectedBeforeAnythingIsOpened() { + TempPackage package(E2bFixture(), {"gemma-4-E2B-it-Q8_0.gguf"}); + FakeRuntimeScope runtime; + OpenAudit audit; + Request request(package, /*context_length=*/20480); + RequireContains(RequireThrows([&] { (void)Build(request.context); }), + "16384"); + TEST_REQUIRE(audit.names.empty()); + TEST_REQUIRE(runtime.nothing_created()); +} + void TestAZeroContextIsRejectedBeforeAnythingIsOpened() { TempPackage package(E2bFixture(), {"gemma-4-E2B-it-Q8_0.gguf"}); FakeRuntimeScope runtime; @@ -614,7 +634,8 @@ void TestAPackageWithNoGgufIsNamedRatherThanGuessedAt() { /// that a package stating `bos_token` and no `bos_token_id` -- which is /// every real Gemma 4 package -- killed the process in the frontend. void RequireTheSevenOverrides(flm::backend::ModelBackend& backend, - int expected_eos, int expected_bos) { + const std::vector& expected_eos, + int expected_bos) { TEST_REQUIRE(backend.id() == std::string(flm::backend::kRaiBackendId)); TEST_REQUIRE(backend.id() == std::string("rai")); // The provenance line in `flm show`: the library the runtime ACTUALLY @@ -628,7 +649,7 @@ void RequireTheSevenOverrides(flm::backend::ModelBackend& backend, TEST_REQUIRE(!backend.poisoned()); const auto eos = backend.forced_eos_ids(); TEST_REQUIRE(eos.has_value()); - TEST_REQUIRE(*eos == std::vector{expected_eos}); + TEST_REQUIRE(*eos == expected_eos); // The seventh. Without it AutoModel::_shared_setup_tokenizer exits the // process on any package that names `bos_token` without stating an id -- // which is what both shipped Google packages do. @@ -641,7 +662,7 @@ void TestTheSevenOverridesOnASyntheticE2bPackage() { TempPackage package(E2bFixture(), {"gemma-4-E2B-it-Q8_0.gguf"}); FakeRuntimeScope runtime; auto backend = Build(Request(package).context); - RequireTheSevenOverrides(*backend, kRealEosId, kRealBosId); + RequireTheSevenOverrides(*backend, {kRealEosId}, kRealBosId); TEST_REQUIRE(runtime.live_objects() > runtime.baseline_objects); } @@ -652,7 +673,7 @@ void TestTheSevenOverridesOnASyntheticE4bPackage() { TempPackage package(E4bFixture(), {"gemma-4-E4B-it-Q8_0.gguf"}); FakeRuntimeScope runtime; auto backend = Build(Request(package).context); - RequireTheSevenOverrides(*backend, kE4bEosId, kE4bBosId); + RequireTheSevenOverrides(*backend, {kE4bEosId}, kE4bBosId); } void TestTheBackendOwnsTheEngineItReportsOn() { @@ -1113,8 +1134,10 @@ void TestTripwireOneTheAieNextFallbackIsGemma4AndItsFamilyCanLoad() { auto& registry = flm::backend::BackendRegistry::instance(); TEST_REQUIRE(registry.has(family, kRai)); const auto traits = registry.traits(family, kRai); - TEST_REQUIRE(traits.max_context_length == - info.at("default_context_length").get()); + const auto default_context = + info.at("default_context_length").get(); + TEST_REQUIRE(default_context <= traits.max_context_length); + TEST_REQUIRE(flm::gemma4::ContextFits(default_context)); } void TestTripwireOneAnUnresolvableTagStillLoadsRatherThanThrowing() { @@ -1137,7 +1160,7 @@ void TestTripwireOneAnUnresolvableTagStillLoadsRatherThanThrowing() { auto backend = flm::backend::BackendRegistry::instance().create( family, kRai, request.context); TEST_REQUIRE(backend->id() == std::string(kRai)); - TEST_REQUIRE(backend->max_decode_length() == 4095); + TEST_REQUIRE(backend->max_decode_length() == context_length - 1); } void TestTripwireOneStxIsUnaffected() { @@ -1367,7 +1390,8 @@ void RequireRealPackageLoads(const std::filesystem::path& gguf, // them. The BOS half is the one that matters most here: Google's // tokenizer_config.json does not state it, so this is the assertion that // says the value the frontend will use really did come off these files. - RequireTheSevenOverrides(*backend, kRealEosId, kRealBosId); + // config.json lists eos_token_id [1, 106], so stops generation too. + RequireTheSevenOverrides(*backend, {kRealEosId, kRealConfigEosId}, kRealBosId); auto* engine = dynamic_cast(&backend->engine()); TEST_REQUIRE(engine != nullptr); TEST_REQUIRE(engine->weight_slot_count() == expected_slots); @@ -1424,9 +1448,11 @@ int main() { RUN_TEST(TestTraitsAreReadableWithoutConstructingAnEngine); RUN_TEST(TestTraitsComeBackOutOfTheRegistryUnchanged); RUN_TEST(TestTheTraitsCeilingIsTheEnginesOwnLimit); + RUN_TEST(TestContextFitsStopsAtTheValidated16384); RUN_TEST(TestPreemptionIsRejectedBeforeAnythingIsOpened); RUN_TEST(TestAnOverLongContextIsRejectedBeforeAnythingIsOpened); + RUN_TEST(TestAContextPastThePagedCeilingIsRejectedBeforeAnythingIsOpened); RUN_TEST(TestAZeroContextIsRejectedBeforeAnythingIsOpened); RUN_TEST(TestAMismatchedConfigJsonIsRejectedAndNoDeviceObjectIsCreated); RUN_TEST(TestABadPackageIsDiagnosedBeforeTheRuntimeIsAcquired); diff --git a/src/test/gemma4_rai/test_gemma4_engine.cpp b/src/test/gemma4_rai/test_gemma4_engine.cpp index 432aae177..55cea9c5d 100644 --- a/src/test/gemma4_rai/test_gemma4_engine.cpp +++ b/src/test/gemma4_rai/test_gemma4_engine.cpp @@ -153,7 +153,8 @@ struct Harness { Gemma4Config shape; explicit Harness(const Fixture& fixture, - const std::function& adjust = {}) { + const std::function& adjust = {}, + std::uint32_t maximum = 4096) { fake_corelib::Reset(); runtime = CorelibRuntime::CreateForTest( CorelibApi::ResolveForTest(fake_corelib::Resolver())); @@ -162,7 +163,7 @@ struct Harness { shape = package->Config(); if (adjust) adjust(shape); engine = gemma4_rai::CreateWithShapeForTest(package, runtime, - LM_Config{}, shape); + LM_Config{}, shape, maximum); } catch (...) { // A CONSTRUCTOR THAT THROWS NEVER RUNS ITS DESTRUCTOR, so without // this the process-wide runtime survives the failed test and @@ -206,6 +207,15 @@ void TestLoadWeightsThrowsAndNamesGguf() { RequireContains(error, "Gemma 4"); } +void TestStreamEnablesTheGemma4BundleBeforeAnyKernelQuery() { + Harness h(E2bFixture()); + const auto& state = fake_corelib::GetState(); + TEST_REQUIRE(state.kernel_bundles.size() == 1); + TEST_REQUIRE(state.kernel_bundles[0].model == "gemma4"); + TEST_REQUIRE(state.kernel_bundles[0].enum_calls_before == 0); + TEST_REQUIRE(!state.enum_calls.empty()); +} + void TestAFreshEngineIsNotPoisonedAndStartsAtPositionZero() { Harness harness(E2bFixture()); TEST_REQUIRE(!harness.engine->poisoned()); @@ -1162,7 +1172,7 @@ void TestPleAndRmsNormAreRecordedAsDispatchesInSequence() { const auto window = [&](void* parent, std::int64_t rows, std::int64_t columns) { const std::int64_t shape[2] = {rows, columns}; void* object = nullptr; - TEST_REQUIRE(functions.create_tensor_window(parent, shape, 2, 0, &object) == + TEST_REQUIRE(functions.create_tensor_window(parent, shape, 2, 0, nullptr, 0, &object) == ryzenai_corelib_status_success); return object; }; @@ -2283,6 +2293,18 @@ void RequireVNormWritesTheCache(const Fixture& fixture) { void TestE2bVNormWritesTheCacheOncePerKvHead() { RequireVNormWritesTheCache(E2bFixture()); } void TestE4bVNormWritesTheCacheOncePerKvHead() { RequireVNormWritesTheCache(E4bFixture()); } +void TestE4bVProjectionLandsThroughAStridedWindow() { + Harness h(E4bFixture()); + (void)h.engine->forward(1); + const auto& windows = fake_corelib::GetState().tensor_windows; + const auto v = std::find_if(windows.begin(), windows.end(), + [](const auto& w) { return w.shape.size() == 3; }); + TEST_REQUIRE(v != windows.end()); + TEST_REQUIRE(v->strides.size() == 3 && v->strides[2] == 1); + TEST_REQUIRE(v->strides[1] == v->shape[2]); + TEST_REQUIRE(v->strides[0] == 4096 * v->strides[1]); +} + /// \brief every attention dispatch gets its own geometry's rotary pair, its /// owner's caches, and the descriptor for its cache role /// @@ -2648,6 +2670,78 @@ void TestAMultiRowContinuationReplaysFromPositionZero() { TEST_REQUIRE(harness.engine->get_current_context_length() == 8); } +/// \brief (position, rows) of every paged_mha dispatch from index `from` on +std::vector> PagedAttentionSince(std::size_t from) { + const auto& dispatches = fake_corelib::GetState().dispatches; + std::vector> passes; + for (std::size_t i = from; i < dispatches.size(); ++i) { + if (dispatches[i].kind == "paged_mha") + passes.emplace_back(dispatches[i].mha_position, dispatches[i].rows); + } + return passes; +} + +/// \brief a paged continuation on a page boundary runs only its own rows +/// \note The server sends a long prompt `max_prefill_len` (4096) tokens at a +/// time, so every call after the first lands on a page boundary. +/// Replaying from position 0 there made each call cost the whole +/// prompt so far. +void TestAPagedContinuationOnAPageBoundaryRunsOnlyItsOwnPage() { + Harness harness(E2bFixture(), {}, 8192); + const auto layers = static_cast(harness.shape.layers); + std::vector first(4096, 5); + (void)harness.engine->prefill(first); + TEST_REQUIRE(harness.engine->get_current_context_length() == 4096); + const auto page0 = PagedAttentionSince(0); + TEST_REQUIRE(page0.size() == layers); + for (const auto& [position, rows] : page0) TEST_REQUIRE(position == 0 && rows == 4096); + + const auto before = fake_corelib::GetState().dispatches.size(); + std::vector second(4096, 7); + (void)harness.engine->prefill(second); + TEST_REQUIRE(harness.engine->get_current_context_length() == 8192); + TEST_REQUIRE(!harness.engine->poisoned()); + const auto page1 = PagedAttentionSince(before); + TEST_REQUIRE(page1.size() == layers); + for (const auto& [position, rows] : page1) TEST_REQUIRE(position == 4096 && rows == 4096); +} + +/// \brief off a page boundary a paged continuation re-runs its page, and +/// inside page 0 that is the whole conversation +void TestAPagedContinuationOffAPageBoundaryReplaysFromItsPage() { + Harness harness(E2bFixture(), {}, 8192); + const auto layers = static_cast(harness.shape.layers); + + std::vector short_prompt{2, 5, 9}; + (void)harness.engine->prefill(short_prompt); + (void)harness.engine->forward(17); + auto before = fake_corelib::GetState().dispatches.size(); + std::vector more{11, 13}; + (void)harness.engine->prefill(more); + TEST_REQUIRE(harness.engine->get_current_context_length() == 6); + const auto in_page0 = PagedAttentionSince(before); + TEST_REQUIRE(in_page0.size() == layers); + for (const auto& [position, rows] : in_page0) TEST_REQUIRE(position == 0 && rows == 128); + + harness.engine->clear_context(); + std::vector long_prompt(4100, 5); + (void)harness.engine->prefill(long_prompt); + TEST_REQUIRE(harness.engine->get_current_context_length() == 4100); + before = fake_corelib::GetState().dispatches.size(); + (void)harness.engine->prefill(more); + TEST_REQUIRE(harness.engine->get_current_context_length() == 4102); + TEST_REQUIRE(!harness.engine->poisoned()); + const auto in_page1 = PagedAttentionSince(before); + TEST_REQUIRE(in_page1.size() == layers); + for (const auto& [position, rows] : in_page1) TEST_REQUIRE(position == 4096 && rows == 4096); + + // And the rows it re-ran are the engine's own: one more continuation + // still finds every token behind it. + std::vector last{19, 23, 29}; + (void)harness.engine->prefill(last); + TEST_REQUIRE(harness.engine->get_current_context_length() == 4105); +} + // --------------------------------------------------------------------------- // The weight cache. @@ -3472,7 +3566,7 @@ void TestTheHostRejectionsAreRaisedOutsideTheTry() { "Gemma 4 request contains no token IDs", "outside the vocabulary of ", "Gemma 4 request exceeds configured context capacity", - "ships no chunked prefill", + "a multi-row pass must start at position 0", "RequireOneRowExtent(padded);", }; for (const auto rejection : kRejections) { @@ -3656,11 +3750,30 @@ void TestAPoisonedEngineStillTearsDownInOrder() { tensors); TEST_REQUIRE(harness.runtime->api()->live_object_count() == 0); } + +void TestAContextPastTheGgufsOwnIsRefused() { + ScopedWeightCache no_cache{"0"}; + RuntimeScope teardown; + fake_corelib::Reset(); + auto runtime = CorelibRuntime::CreateForTest( + CorelibApi::ResolveForTest(fake_corelib::Resolver())); + auto package = Gemma4GgufPackage::Open(E2bFixture().path); + auto shape = package->Config(); + shape.context_length = 8192; + const auto error = RequireThrows([&] { + (void)gemma4_rai::CreateWithShapeForTest(package, runtime, LM_Config{}, shape, 12288); + }); + RequireContains(error, "past this GGUF's context of 8192"); + package.reset(); + runtime.reset(); + CorelibRuntime::ShutdownProcess(); +} } // namespace int main() { #define RUN_TEST(name) RunTest(&name, #name) RUN_TEST(TestLoadWeightsThrowsAndNamesGguf); + RUN_TEST(TestStreamEnablesTheGemma4BundleBeforeAnyKernelQuery); RUN_TEST(TestAFreshEngineIsNotPoisonedAndStartsAtPositionZero); RUN_TEST(TestE2bHoldsTenWeightsPerOwningLayerAndSixPerSharingOne); RUN_TEST(TestE4bHoldsTenWeightsPerOwningLayerAndSixPerSharingOne); @@ -3691,6 +3804,7 @@ int main() { RUN_TEST(TestE4bEveryLayerBindsItsOwnGeometrysWidths); RUN_TEST(TestE2bVNormWritesTheCacheOncePerKvHead); RUN_TEST(TestE4bVNormWritesTheCacheOncePerKvHead); + RUN_TEST(TestE4bVProjectionLandsThroughAStridedWindow); RUN_TEST(TestE2bAttentionBindsItsOwnRotaryPairAndItsOwnersCaches); RUN_TEST(TestE4bAttentionBindsItsOwnRotaryPairAndItsOwnersCaches); RUN_TEST(TestTheReturnedLogitsAreSoftcapped); @@ -3700,6 +3814,8 @@ int main() { RUN_TEST(TestACacheReadResolvesThroughTheOwner); RUN_TEST(TestTheHostRejectionsFireBeforeAnyDispatch); RUN_TEST(TestAMultiRowContinuationReplaysFromPositionZero); + RUN_TEST(TestAPagedContinuationOnAPageBoundaryRunsOnlyItsOwnPage); + RUN_TEST(TestAPagedContinuationOffAPageBoundaryReplaysFromItsPage); RUN_TEST(TestE2bSecondLoadOfTheSameGgufBindsFromTheCache); RUN_TEST(TestE4bSecondLoadOfTheSameGgufBindsFromTheCache); RUN_TEST(TestACacheWrittenForE2bIsNotBoundForE4b); @@ -3723,5 +3839,6 @@ int main() { RUN_TEST(TestE4bARefusedFirstDispatchDoesNotPoison); RUN_TEST(TestTheHostRejectionsAreRaisedOutsideTheTry); RUN_TEST(TestAPoisonedEngineStillTearsDownInOrder); + RUN_TEST(TestAContextPastTheGgufsOwnIsRefused); #undef RUN_TEST } diff --git a/src/test/gemma4_rai/test_gemma4_gguf.cpp b/src/test/gemma4_rai/test_gemma4_gguf.cpp index b87e16f7a..918c36749 100644 --- a/src/test/gemma4_rai/test_gemma4_gguf.cpp +++ b/src/test/gemma4_rai/test_gemma4_gguf.cpp @@ -25,6 +25,7 @@ void TestConfigFromMetadataReadsE2b() { TEST_REQUIRE(cfg.kv_heads == 1); TEST_REQUIRE(cfg.kv_layers == 15); TEST_REQUIRE(cfg.sliding_window == 512); + TEST_REQUIRE(cfg.context_length == 131072); TEST_REQUIRE(cfg.vocab == 262144); TEST_REQUIRE(cfg.rope_theta == 1000000.0); TEST_REQUIRE(cfg.rope_theta_swa == 10000.0); @@ -52,6 +53,21 @@ void TestConfigFromMetadataReadsE4b() { TEST_REQUIRE(cfg.ple_dim == 256); } +void TestContextLengthIsReadFromTheFile() { + auto builder = gemma4_fixture::MakeBuilder(gemma4_fixture::E2bOptions()); + builder.SetMetadata("gemma4.context_length", std::uint32_t{8192}); + auto file = GgufFile::Open(builder.Write()); + TEST_REQUIRE(ConfigFromMetadata(*file).context_length == 8192); +} + +void TestAFileWithoutAContextLengthGetsOnePage() { + auto builder = gemma4_fixture::MakeBuilder(gemma4_fixture::E2bOptions()); + builder.RemoveMetadata("gemma4.context_length"); + auto file = GgufFile::Open(builder.Write()); + TEST_REQUIRE(!file->HasMetadata("gemma4.context_length")); + TEST_REQUIRE(ConfigFromMetadata(*file).context_length == 4096); +} + // THE ONE THAT MATTERS. gemma4.attention.shared_kv_layers is the file's // SHARING count (20 on E2B, 18 on E4B), never the OWNING count. Pinning // cfg.kv_layers against the real metadata values here means the @@ -322,6 +338,8 @@ int main() { #define RUN_TEST(name) RunTest(&name, #name) RUN_TEST(TestConfigFromMetadataReadsE2b); RUN_TEST(TestConfigFromMetadataReadsE4b); + RUN_TEST(TestContextLengthIsReadFromTheFile); + RUN_TEST(TestAFileWithoutAContextLengthGetsOnePage); RUN_TEST(TestKvLayersIsBlockCountMinusTheFilesSharingCount); RUN_TEST(TestHeadDimKeysAreNotReadTheObviousWayRound); RUN_TEST(TestTensorCountIsSixPlusSeventeenPerLayer); diff --git a/src/test/gemma4_rai/test_gemma4_shape_plan.cpp b/src/test/gemma4_rai/test_gemma4_shape_plan.cpp index a5d70816d..c1c5b08eb 100644 --- a/src/test/gemma4_rai/test_gemma4_shape_plan.cpp +++ b/src/test/gemma4_rai/test_gemma4_shape_plan.cpp @@ -586,6 +586,50 @@ void TestPleRowsRoundUpToAShippedBucket() { /// few columns wider is a different artifact. The plan must not take it as /// covering the logical shape. The N side and the K side are broken /// independently, so a fix that only compared one would not pass this. +/// \brief the paged kernels corelib 0.26's gemma4 bundle ships: the decode +/// step and the page-sized chunk, nothing in between +void ShipPagedChunkOnly(const Gemma4Config& config) { + auto& state = fake_corelib::GetState(); + for (const auto head : {config.head_dim, config.global_head_dim}) { + for (const std::int64_t rows : {std::int64_t{1}, flm::gemma4::kPageSize}) { + state.shipped_paged.push_back( + {config.q_heads, config.kv_heads, rows, head, flm::gemma4::kPageSize}); + } + } +} + +void TestAPagedPlanSizesAttentionFromFlatMhaAtPositionZero() { + fake_corelib::Reset(); + auto api = Api(); + auto* stream = fake_corelib::MakeStreamForTest(); + const auto config = E2bConfigForTest(); + ShipFullGrid(config); + ShipPagedChunkOnly(config); + const auto plan = Gemma4ShapePlan::Build(api, stream, config, 2 * flm::gemma4::kPageSize); + TEST_REQUIRE(plan.ForRows(19, /*swa=*/true).flat_mha_rows == 64); + TEST_REQUIRE(plan.ForRows(100, /*swa=*/false).flat_mha_rows == 128); + TEST_REQUIRE(plan.ForRows(4096, /*swa=*/true).flat_mha_rows == 4096); +} + +void TestAPagedPlanRefusesAGridWithoutThePageSizedChunk() { + fake_corelib::Reset(); + auto api = Api(); + auto* stream = fake_corelib::MakeStreamForTest(); + const auto config = E2bConfigForTest(); + ShipFullGrid(config); + for (const auto head : {config.head_dim, config.global_head_dim}) { + fake_corelib::GetState().shipped_paged.push_back( + {config.q_heads, config.kv_heads, 1, head, flm::gemma4::kPageSize}); + } + std::string error; + try { + (void)Gemma4ShapePlan::Build(api, stream, config, 2 * flm::gemma4::kPageSize); + } catch (const std::runtime_error& e) { + error = e.what(); + } + RequireContains(error, "paged_mha chunk has no kernel"); +} + void TestPlanRejectsAGridWithoutItsExactKAndN() { const auto config = E2bConfigForTest(); RequireGridBreaks(config, [&config](auto& state) { @@ -619,13 +663,13 @@ void TestFakeRejectsAWindowThatDoesNotFitItsParent() { // 16 elements starting 16 in: exactly reaches the end, so it is legal. const std::int64_t window_shape[] = {2, 8}; ryzenai_corelib_tensor_ptr fits = nullptr; - TEST_REQUIRE(fns.create_tensor_window(parent, window_shape, 2, 16, &fits) == + TEST_REQUIRE(fns.create_tensor_window(parent, window_shape, 2, 16, nullptr, 0, &fits) == ryzenai_corelib_status_success); TEST_REQUIRE(fits != nullptr); // The same shape 8 elements further along runs 8 past the end. ryzenai_corelib_tensor_ptr overruns = nullptr; - TEST_REQUIRE(fns.create_tensor_window(parent, window_shape, 2, 24, &overruns) != + TEST_REQUIRE(fns.create_tensor_window(parent, window_shape, 2, 24, nullptr, 0, &overruns) != ryzenai_corelib_status_success); TEST_REQUIRE(overruns == nullptr); @@ -717,9 +761,9 @@ void TestFakeRejectsSsmlpOperandsThatAliasOneAllocation() { ryzenai_corelib_status_success); ryzenai_corelib_tensor_ptr front = nullptr; ryzenai_corelib_tensor_ptr back = nullptr; - TEST_REQUIRE(fns.create_tensor_window(parent, shape, 2, 0, &front) == + TEST_REQUIRE(fns.create_tensor_window(parent, shape, 2, 0, nullptr, 0, &front) == ryzenai_corelib_status_success); - TEST_REQUIRE(fns.create_tensor_window(parent, shape, 2, 64 * 1536, &back) == + TEST_REQUIRE(fns.create_tensor_window(parent, shape, 2, 64 * 1536, nullptr, 0, &back) == ryzenai_corelib_status_success); // THE POINT: distinct handles, one allocation. TEST_REQUIRE(front != back); @@ -834,6 +878,8 @@ int main() { RUN_TEST(TestDispatchRecordsCarryTheDescriptorTheyRanWith); RUN_TEST(TestE4bGroupDirectionAndLmHeadPlan); RUN_TEST(TestPleRowsRoundUpToAShippedBucket); + RUN_TEST(TestAPagedPlanSizesAttentionFromFlatMhaAtPositionZero); + RUN_TEST(TestAPagedPlanRefusesAGridWithoutThePageSizedChunk); RUN_TEST(TestPlanRejectsAGridWithoutItsExactKAndN); RUN_TEST(TestFakeRejectsAWindowThatDoesNotFitItsParent); RUN_TEST(TestFakeBoundsChecksTensorWriteAndRead); diff --git a/src/test/phi4_rai/fake_corelib.cpp b/src/test/phi4_rai/fake_corelib.cpp index 9658413a7..1eb30bc8a 100644 --- a/src/test/phi4_rai/fake_corelib.cpp +++ b/src/test/phi4_rai/fake_corelib.cpp @@ -29,6 +29,7 @@ struct FakeObject { std::vector shape; std::size_t byte_size{}; std::size_t window_offset{}; + std::vector strides; std::shared_ptr storage; }; @@ -198,6 +199,52 @@ struct TypedFake { auto* out = std::get<2>(arguments); if (out) *out = status == ryzenai_corelib_status_success ? NewObject("stream") : nullptr; return status; + } else if constexpr (std::is_same_v) { + const char* model = std::get<1>(arguments); + state.kernel_bundles.emplace_back(model ? model : ""); + return Status(Tag::name); + } else if constexpr (std::is_same_v) { + const auto status = Status(Tag::name); + auto* out = std::get<2>(arguments); + if (out) *out = status == ryzenai_corelib_status_success ? NewObject("model_weights") : nullptr; + return status; + } else if constexpr (std::is_same_v) { + const auto status = Status(Tag::name); + const auto* options = std::get<2>(arguments); + auto* out = std::get<3>(arguments); + state.model_max_sequence_length = options ? options->max_sequence_length : 0; + ++state.model_creates; + if (out) *out = status == ryzenai_corelib_status_success ? NewObject("model") : nullptr; + return status; + } else if constexpr (std::is_same_v) { + auto* out = std::get<1>(arguments); + if (out) { + constexpr std::int64_t kPage = 4096; + const auto wanted = std::max(state.model_max_sequence_length, 1); + *out = {state.model_vocab, (wanted + kPage - 1) / kPage * kPage, 16, 4}; + } + return Status(Tag::name); + } else if constexpr (std::is_same_v) { + const auto* options = std::get<5>(arguments); + state.model_forwards.push_back( + {std::get<3>(arguments), + options ? options->position : RYZENAI_CORELIB_POSITION_CONTINUE, + options && options->logits}); + if (options && options->logits) + std::fill_n(options->logits, + std::min(options->logits_capacity, + static_cast(state.model_vocab)), + 0.5f); + return Status(Tag::name); + } else if constexpr (std::is_same_v) { + const auto status = Status(Tag::name); + ++state.state_saves; + auto* out = std::get<1>(arguments); + if (out) *out = status == ryzenai_corelib_status_success ? NewObject("snapshot") : nullptr; + return status; + } else if constexpr (std::is_same_v) { + ++state.state_restores; + return Status(Tag::name); } else if constexpr (std::is_same_v) { const auto status = Status(Tag::name); const auto type = std::get<0>(arguments); @@ -222,7 +269,7 @@ struct TypedFake { const auto* shape = std::get<1>(arguments); const auto shape_len = std::get<2>(arguments); const auto offset = std::get<3>(arguments); - auto* out = std::get<4>(arguments); + auto* out = std::get<6>(arguments); if (out) *out = nullptr; if (status == ryzenai_corelib_status_success && out && shape) { auto* object = static_cast(NewObject("window")); @@ -234,8 +281,11 @@ struct TypedFake { } object->shape.assign(shape, shape + shape_len); object->byte_size = Elements(object->shape) * TypeBytes(object->data_type); + const auto* strides = std::get<4>(arguments); + const auto strides_len = std::get<5>(arguments); + if (strides) object->strides.assign(strides, strides + strides_len); *out = object; - state.tensor_windows.push_back({parent, object->shape, offset, object}); + state.tensor_windows.push_back({parent, object->shape, offset, object, object->strides}); } return status; } else if constexpr (std::is_same_v) { @@ -522,6 +572,21 @@ struct TypedFake { } const auto status = Status(Tag::name); if (status != ryzenai_corelib_status_success) return status; + if constexpr (std::is_same_v) { + // Like corelib: a 3-D output is the cache-writing form and must + // be a strided window [k * head, head, 1] with k >= its rows. + const auto* output = static_cast(std::get<3>(arguments)); + if (output && output->shape.size() == 3) { + const auto& strides = output->strides; + const auto head = output->shape[2]; + if (head <= 0 || strides.size() != 3 || strides[2] != 1 || strides[1] != head || + strides[0] % head != 0 || strides[0] / head < output->shape[1]) { + state.detail = "a cache-writing 3-D output must be a window with strides " + "[max_seq * head_size, head_size, 1]"; + return ryzenai_corelib_status_bad_argument; + } + } + } if constexpr (std::is_same_v) { // Real corelib reads the rotary tables on the host and rejects // anything but a host view for them. @@ -649,10 +714,17 @@ void Reset() { state.dispatches.clear(); state.tensor_writes.clear(); state.call_log.clear(); + state.kernel_bundles.clear(); state.active_weight_creates = 0; state.maximum_active_weight_creates = 0; state.work_in_flight = false; state.fail_after_submit.clear(); + state.model_max_sequence_length = -1; + state.model_vocab = 8; + state.model_creates = 0; + state.state_saves = 0; + state.state_restores = 0; + state.model_forwards.clear(); } flm::corelib::CorelibApi::Resolver Resolver() { diff --git a/src/test/phi4_rai/fake_corelib.hpp b/src/test/phi4_rai/fake_corelib.hpp index 631591239..6c4939450 100644 --- a/src/test/phi4_rai/fake_corelib.hpp +++ b/src/test/phi4_rai/fake_corelib.hpp @@ -44,6 +44,7 @@ struct TensorWindowRecord { std::vector shape; std::size_t offset; void* object; + std::vector strides; }; struct WeightCreateRecord { @@ -87,6 +88,12 @@ struct WeightFromFileRecord { std::uint64_t size{}; }; +struct ModelForwardRecord { + std::size_t rows{}; + std::size_t position{}; + bool logits{}; +}; + struct State { /// \note Derived from the vendor header, not written out: the runtime /// check in corelib_api.cpp compares against the same macros, so a @@ -135,10 +142,17 @@ struct State { std::vector dispatches; std::vector tensor_writes; std::vector call_log; + std::vector kernel_bundles; std::atomic active_weight_creates{0}; std::atomic maximum_active_weight_creates{0}; bool work_in_flight{false}; std::string fail_after_submit; + std::int64_t model_max_sequence_length{-1}; + std::int64_t model_vocab{8}; + std::size_t model_creates{}; + std::size_t state_saves{}; + std::size_t state_restores{}; + std::vector model_forwards; }; State& GetState(); diff --git a/src/test/phi4_rai/run_real_rai_acceptance.ps1 b/src/test/phi4_rai/run_real_rai_acceptance.ps1 index a553d5283..21c02d4ab 100644 --- a/src/test/phi4_rai/run_real_rai_acceptance.ps1 +++ b/src/test/phi4_rai/run_real_rai_acceptance.ps1 @@ -1,7 +1,7 @@ param( [string]$FlmExe = 'src/build-rai/Release/flm.exe', [string]$Model = 'phi4-mini-it:4b', - [string]$CorelibDll = 'C:/Users/chiz/work/ryzenai-corelib/install/bin/ryzenai_corelib.dll', + [string]$CorelibDll = 'C:/Users/chiz/work/ryzenai-corelib-0.25/install/bin/ryzenai_corelib.dll', [string]$Output = 'src/build-rai/phi4-gguf-rai-acceptance.json', [int]$Port = 52625, [string]$Python = 'python', @@ -137,7 +137,7 @@ function CurlStream([string]$name,[string]$path,$body){$f=BodyFile $name $body;$ function CurlBackground([string]$name,[string]$path,$body,[string]$outFile){$f=BodyFile $name $body;return (Start-Process curl.exe -ArgumentList @('-sS','-N','-H',$ContentTypeArg,'-d',"@$f","http://127.0.0.1:$Port$path") -RedirectStandardOutput $outFile -PassThru)} try{ $record.host.computer=$env:COMPUTERNAME;$record.host.cpu=(Get-CimInstance Win32_Processor).Name;$record.host.npu=(Get-CimInstance Win32_PnPEntity|Where-Object Name -match 'NPU|Neural').Name;$os=Get-CimInstance Win32_OperatingSystem;$record.host.windows="$($os.Caption) $($os.Version) build $($os.BuildNumber)";$record.host.power=(powercfg /getactivescheme|Out-String).Trim() - $gitRev={param($d) if(-not (Test-Path $d)){return 'unknown'};$keep=$ErrorActionPreference;$ErrorActionPreference='SilentlyContinue';$r=$null;try{$r=(& git -C $d rev-parse HEAD 2>$null | Select-Object -First 1)}catch{$r=$null};$ErrorActionPreference=$keep;$global:LASTEXITCODE=0;if([string]::IsNullOrWhiteSpace($r)){'unknown'}else{([string]$r).Trim()}};$record.provenance.fastflow=(& $gitRev $root);$coreRoot=(Resolve-Path (Join-Path $runtime '..')).Path;$record.provenance.corelib=(& $gitRev $coreRoot);$record.provenance.corelib_abi='0.5.0';$record.provenance.gguf_revision='78eb92a46fc37e6b524df991ed9aca9bc6aa7b80';$record.provenance.tokenizer_revision='cfbefacb99257ffa30c83adab238a50856ac3083';$record.provenance.corelib_sha256=(Get-FileHash $core -Algorithm SHA256).Hash.ToLower() + $gitRev={param($d) if(-not (Test-Path $d)){return 'unknown'};$keep=$ErrorActionPreference;$ErrorActionPreference='SilentlyContinue';$r=$null;try{$r=(& git -C $d rev-parse HEAD 2>$null | Select-Object -First 1)}catch{$r=$null};$ErrorActionPreference=$keep;$global:LASTEXITCODE=0;if([string]::IsNullOrWhiteSpace($r)){'unknown'}else{([string]$r).Trim()}};$record.provenance.fastflow=(& $gitRev $root);$coreRoot=(Resolve-Path (Join-Path $runtime '..')).Path;$record.provenance.corelib=(& $gitRev $coreRoot);$record.provenance.corelib_abi='0.25.0';$record.provenance.gguf_revision='78eb92a46fc37e6b524df991ed9aca9bc6aa7b80';$record.provenance.tokenizer_revision='cfbefacb99257ffa30c83adab238a50856ac3083';$record.provenance.corelib_sha256=(Get-FileHash $core -Algorithm SHA256).Hash.ToLower() Cmd "$exe check $Model" {&$exe check $Model|Out-Host;if($LASTEXITCODE-ne 0){throw 'check failed'}} # The packed-weight cache is written into the model directory, so a second run # finds two more files there. They are derived, not downloaded: name them and diff --git a/src/test/phi4_rai/test_phi4_engine.cpp b/src/test/phi4_rai/test_phi4_engine.cpp index 50f87242a..90e07f6f5 100644 --- a/src/test/phi4_rai/test_phi4_engine.cpp +++ b/src/test/phi4_rai/test_phi4_engine.cpp @@ -85,6 +85,18 @@ void TestEngineCreatesOneStreamAndPersistentHelperSizedTensors() { TEST_REQUIRE(state.tensor_creates[7].shape == std::vector({1, 200064})); } +void TestStreamEnablesThePhi4MiniBundleBeforeAnyKernelQuery() { + Harness h; + const auto& state = fake_corelib::GetState(); + TEST_REQUIRE(state.kernel_bundles == std::vector({"phi4mini"})); + const auto& log = state.call_log; + const auto created = std::find(log.begin(), log.end(), "ryzenai_corelib_create_stream"); + const auto enabled = std::find(log.begin(), log.end(), "ryzenai_corelib_stream_enable_kernel_bundle"); + const auto queried = std::find(log.begin(), log.end(), "ryzenai_corelib_matmul_bf16_enum_kernels"); + TEST_REQUIRE(queried != log.end()); + TEST_REQUIRE(created < enabled && enabled < queried); +} + void TestRopeTablesAreHostViewsNotDeviceTensors() { Harness h; const auto& views = fake_corelib::GetState().host_view_creates; @@ -431,6 +443,16 @@ void TestVProjectionWritesWindowAtPositionTimes128() { TEST_REQUIRE(fake_corelib::GetState().dispatches[2].window_offset == 7 * 128); } +void TestTheVProjectionWritesTheCacheThroughAStridedWindow() { + Harness h; + (void)h.engine->forward(1); + const auto& windows = fake_corelib::GetState().tensor_windows; + const auto v = std::find_if(windows.begin(), windows.end(), + [](const auto& w) { return w.shape.size() == 3; }); + TEST_REQUIRE(v != windows.end()); + TEST_REQUIRE(v->strides == std::vector({4096 * 128, 128, 1})); +} + void TestEachLayerOrdersQKVThenMhaThenOThenSsmlpOnOneStream() { Harness h; (void)h.engine->forward(1); @@ -586,7 +608,7 @@ void TestFakeTensorWindowRetainsAndPropagatesParentStorage() { const std::array window_shape{4}; void* window = nullptr; api->Check(api->functions().create_tensor_window( - parent, window_shape.data(), window_shape.size(), 4, &window), "create window"); + parent, window_shape.data(), window_shape.size(), 4, nullptr, 0, &window), "create window"); std::array read{}; api->Check(api->functions().tensor_read(window, ryzenai_corelib_data_type_bf16, read.data(), read.size(), 0), "read window"); @@ -726,6 +748,7 @@ void TestTenSequentialLoadsReleaseEveryObjectAndNeverEmitAllZeroLogits() { int main() { #define RUN_TEST(name) RunTest(&name, #name) RUN_TEST(TestEngineCreatesOneStreamAndPersistentHelperSizedTensors); + RUN_TEST(TestStreamEnablesThePhi4MiniBundleBeforeAnyKernelQuery); RUN_TEST(TestRopeTablesAreHostViewsNotDeviceTensors); RUN_TEST(TestEngineAllocatesMaximaAcrossAllRowsAndConsumers); RUN_TEST(TestEngineCreatesExactly129MatmulAnd32SsmlpWeights); @@ -742,6 +765,7 @@ int main() { RUN_TEST(TestPrefillDecodesEmbeddingRowsAndAdvancesPosition); RUN_TEST(TestDecodeUsesOneRowAndAdvancesPosition); RUN_TEST(TestVProjectionWritesWindowAtPositionTimes128); + RUN_TEST(TestTheVProjectionWritesTheCacheThroughAStridedWindow); RUN_TEST(TestEachLayerOrdersQKVThenMhaThenOThenSsmlpOnOneStream); RUN_TEST(TestPrefillStagesTheSameFp32EmbeddingIntoHiddenAndResidual); RUN_TEST(TestBuffersAreZeroPaddedBeforeSubmissionForEachRowBucket); diff --git a/src/test/qwen36_rai/CMakeLists.txt b/src/test/qwen36_rai/CMakeLists.txt new file mode 100644 index 000000000..17d2f316e --- /dev/null +++ b/src/test/qwen36_rai/CMakeLists.txt @@ -0,0 +1,42 @@ +cmake_minimum_required(VERSION 3.22) +project(qwen36_rai_tests LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +if(NOT WIN32) + message(FATAL_ERROR "Qwen3.6 corelib tests currently require Windows") +endif() + +find_path(RYZENAI_CORELIB_INCLUDE_DIR NAMES ryzenai/corelib.h REQUIRED) +find_path(BOOST_INCLUDE_DIR NAMES boost/program_options.hpp + HINTS "$ENV{CONDA_PREFIX}/Library/include" + "$ENV{USERPROFILE}/anaconda3/Library/include" REQUIRED) +find_path(XRT_INCLUDE_DIR NAMES xrt/xrt_bo.h + HINTS "$ENV{XRT_INCLUDE_DIR}" + "C:/dev/XRT/src/runtime_src/core/include" REQUIRED) +set(FLM_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../..") + +add_executable(test_qwen36_rai + test_qwen36_rai.cpp + "${CMAKE_CURRENT_LIST_DIR}/../phi4_rai/fake_corelib.cpp" + "${FLM_SOURCE_DIR}/common/rai/corelib_api.cpp" + "${FLM_SOURCE_DIR}/common/rai/corelib_runtime.cpp" + "${FLM_SOURCE_DIR}/common/rai/host_ops.cpp" + "${FLM_SOURCE_DIR}/common/models/qwen36/rai/aie_next/qwen36_rai.cpp") +target_include_directories(test_qwen36_rai PRIVATE + "${CMAKE_CURRENT_LIST_DIR}" + "${FLM_SOURCE_DIR}/include" + "${RYZENAI_CORELIB_INCLUDE_DIR}" + "${BOOST_INCLUDE_DIR}" + "${XRT_INCLUDE_DIR}") +target_compile_definitions(test_qwen36_rai PRIVATE + RYZENAI_CORELIB_STATIC=1 USEAVX2=1 DISABLE_ABI_CHECK=1 + _ENABLE_EXTENDED_ALIGNED_STORAGE WIN32_LEAN_AND_MEAN NOMINMAX) +target_compile_options(test_qwen36_rai PRIVATE + $<$:/wd4005 /wd4244>) +target_link_directories(test_qwen36_rai PRIVATE "${XRT_INCLUDE_DIR}/../lib") +target_link_libraries(test_qwen36_rai PRIVATE xrt_coreutil) + +include(CTest) +add_test(NAME test_qwen36_rai COMMAND test_qwen36_rai) diff --git a/src/test/qwen36_rai/test_qwen36_rai.cpp b/src/test/qwen36_rai/test_qwen36_rai.cpp new file mode 100644 index 000000000..c046ed665 --- /dev/null +++ b/src/test/qwen36_rai/test_qwen36_rai.cpp @@ -0,0 +1,236 @@ +/// \file test_qwen36_rai.cpp +/// \brief Qwen3.6 on rai, hardware-free: capacity, prefill shape, rewind and +/// snapshots against the fake corelib +#include "models/qwen36/rai/aie_next/qwen36_rai.hpp" +#include "../phi4_rai/fake_corelib.hpp" +#include "../phi4_rai/test_support.hpp" + +#include +#include +#include +#include + +namespace { +using flm::corelib::CorelibApi; +using flm::corelib::CorelibRuntime; +using flm::qwen36::qwen36_rai; + +struct Harness { + std::shared_ptr runtime; + std::unique_ptr engine; + + explicit Harness(std::uint32_t maximum) { + fake_corelib::Reset(); + runtime = CorelibRuntime::CreateForTest(CorelibApi::ResolveForTest(fake_corelib::Resolver())); + try { + engine = std::make_unique("C:/fake/qwen36.gguf", runtime, maximum); + } catch (...) { + Shutdown(); + throw; + } + } + ~Harness() { Shutdown(); } + + void Shutdown() { + engine.reset(); + runtime.reset(); + CorelibRuntime::ShutdownProcess(); + } +}; + +std::vector Ids(std::size_t count) { + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 1); + return ids; +} + +void TestTheModelIsCreatedAtTheConfiguredCapacity() { + Harness h(32768); + TEST_REQUIRE(fake_corelib::GetState().model_max_sequence_length == 32768); +} + +void TestTheModelsOwnLimitIsAccepted() { + Harness h(262144); + TEST_REQUIRE(fake_corelib::GetState().model_max_sequence_length == 262144); +} + +void TestAMaximumPastTheModelsOwnIsRefused() { + fake_corelib::Reset(); + auto runtime = CorelibRuntime::CreateForTest(CorelibApi::ResolveForTest(fake_corelib::Resolver())); + (void)RequireThrows([&] { qwen36_rai engine("C:/fake/qwen36.gguf", runtime, 262145); }); + runtime.reset(); + CorelibRuntime::ShutdownProcess(); +} + +void TestAPromptIsOneForward() { + Harness h(8192); + auto ids = Ids(5000); + (void)h.engine->prefill(ids); + const auto& forwards = fake_corelib::GetState().model_forwards; + TEST_REQUIRE(forwards.size() == 1); + TEST_REQUIRE(forwards[0].rows == 5000 && forwards[0].logits); + TEST_REQUIRE(h.engine->get_current_context_length() == 5000); +} + +void TestRestoreLoadsTheSnapshotWithoutReplaying() { + Harness h(4096); + auto ids = Ids(10); + (void)h.engine->prefill(ids); + TEST_REQUIRE(h.engine->checkpoint() == 10); + (void)h.engine->forward(1); + (void)h.engine->forward(2); + const auto forwards_before = fake_corelib::GetState().model_forwards.size(); + TEST_REQUIRE(h.engine->restore() == 10); + TEST_REQUIRE(fake_corelib::GetState().state_saves == 1); + TEST_REQUIRE(fake_corelib::GetState().state_restores == 1); + TEST_REQUIRE(fake_corelib::GetState().model_forwards.size() == forwards_before); + TEST_REQUIRE(h.engine->get_current_context_length() == 10); +} + +void TestRewindingElsewhereReplaysInOnePass() { + Harness h(4096); + auto ids = Ids(10); + (void)h.engine->prefill(ids); + (void)h.engine->forward(1); + h.engine->set_context_length(5); + const auto& forwards = fake_corelib::GetState().model_forwards; + TEST_REQUIRE(forwards.back().rows == 5 && !forwards.back().logits); + TEST_REQUIRE(h.engine->get_current_context_length() == 5); +} + +void TestRewindingOntoTheCheckpointLoadsItsSnapshot() { + Harness h(4096); + auto ids = Ids(10); + (void)h.engine->prefill(ids); + (void)h.engine->checkpoint(); + (void)h.engine->forward(1); + (void)h.engine->forward(2); + const auto forwards_before = fake_corelib::GetState().model_forwards.size(); + h.engine->set_context_length(10); + TEST_REQUIRE(fake_corelib::GetState().state_restores == 1); + TEST_REQUIRE(fake_corelib::GetState().model_forwards.size() == forwards_before); + TEST_REQUIRE(h.engine->get_current_context_length() == 10); +} + +void TestRewindingOntoADivergedCheckpointReplays() { + Harness h(4096); + auto ids = Ids(10); + (void)h.engine->prefill(ids); + (void)h.engine->checkpoint(); + h.engine->set_context_length(5); + std::vector other{101, 102, 103, 104, 105}; + (void)h.engine->prefill(other); + (void)h.engine->forward(1); + (void)h.engine->forward(2); + auto& state = fake_corelib::GetState(); + const auto resets_before = state.call_counts["ryzenai_corelib_model_reset"]; + const auto forwards_before = state.model_forwards.size(); + h.engine->set_context_length(10); + TEST_REQUIRE(state.state_restores == 0); + TEST_REQUIRE(state.call_counts["ryzenai_corelib_model_reset"] == resets_before + 1); + TEST_REQUIRE(state.model_forwards.size() == forwards_before + 1); + TEST_REQUIRE(state.model_forwards.back().rows == 10 && !state.model_forwards.back().logits); + TEST_REQUIRE(h.engine->get_current_context_length() == 10); +} + +void TestAFailedRestorePoisonsTheEngine() { + Harness h(4096); + auto ids = Ids(10); + (void)h.engine->prefill(ids); + (void)h.engine->checkpoint(); + (void)h.engine->forward(1); + fake_corelib::GetState().statuses["ryzenai_corelib_model_state_restore"] = + ryzenai_corelib_status_failure; + (void)RequireThrows([&] { (void)h.engine->restore(); }); + TEST_REQUIRE(h.engine->poisoned()); + RequireContains(RequireThrows([&] { (void)h.engine->forward(2); }), "poisoned"); +} + +void TestAFailedForwardPoisonsTheEngineAndDropsTheCheckpoint() { + Harness h(4096); + auto ids = Ids(10); + (void)h.engine->prefill(ids); + (void)h.engine->checkpoint(); + auto& state = fake_corelib::GetState(); + state.statuses["ryzenai_corelib_model_forward"] = ryzenai_corelib_status_failure; + (void)RequireThrows([&] { (void)h.engine->forward(1); }); + TEST_REQUIRE(h.engine->poisoned()); + state.statuses.erase("ryzenai_corelib_model_forward"); + const auto forwards = state.model_forwards.size(); + RequireContains(RequireThrows([&] { (void)h.engine->forward(2); }), "poisoned"); + RequireContains(RequireThrows([&] { (void)h.engine->prefill(ids); }), "poisoned"); + RequireContains(RequireThrows([&] { (void)h.engine->restore(); }), "poisoned"); + RequireContains(RequireThrows([&] { (void)h.engine->checkpoint(); }), "poisoned"); + TEST_REQUIRE(state.model_forwards.size() == forwards); + TEST_REQUIRE(state.state_restores == 0); +} + +void TestCheckpointingTheSameStateTwiceSavesOnce() { + Harness h(4096); + auto ids = Ids(10); + (void)h.engine->prefill(ids); + TEST_REQUIRE(h.engine->checkpoint() == 10); + TEST_REQUIRE(h.engine->checkpoint() == 10); + TEST_REQUIRE(fake_corelib::GetState().state_saves == 1); +} + +void TestRestoringRightAfterACheckpointLoadsNothing() { + Harness h(4096); + auto ids = Ids(10); + (void)h.engine->prefill(ids); + (void)h.engine->checkpoint(); + TEST_REQUIRE(h.engine->restore() == 10); + TEST_REQUIRE(fake_corelib::GetState().state_restores == 0); +} + +void TestShrinkingBelowTheCheckpointDropsIt() { + Harness h(16384); + auto ids = Ids(10); + (void)h.engine->prefill(ids); + (void)h.engine->checkpoint(); + h.engine->set_context_length(2); + h.engine->update_max_length(5); + TEST_REQUIRE(h.engine->restore() == -1); + TEST_REQUIRE(fake_corelib::GetState().state_restores == 0); +} + +void TestGrowingTheMaximumRecreatesTheModelAndKeepsTheSequence() { + Harness h(4096); + auto ids = Ids(10); + (void)h.engine->prefill(ids); + h.engine->update_max_length(16384); + const auto& state = fake_corelib::GetState(); + TEST_REQUIRE(state.model_creates == 2); + TEST_REQUIRE(state.model_max_sequence_length == 16384); + TEST_REQUIRE(state.state_saves == 1 && state.state_restores == 1); + TEST_REQUIRE(h.engine->get_current_context_length() == 10); +} + +void TestShrinkingWithinTheCapacityKeepsTheModel() { + Harness h(16384); + h.engine->update_max_length(8192); + TEST_REQUIRE(fake_corelib::GetState().model_creates == 1); +} + +} // namespace + +int main() { + RunTest(TestTheModelIsCreatedAtTheConfiguredCapacity, "the model is created at the configured capacity"); + RunTest(TestTheModelsOwnLimitIsAccepted, "the model's own limit is accepted"); + RunTest(TestAMaximumPastTheModelsOwnIsRefused, "a maximum past the model's own is refused"); + RunTest(TestAPromptIsOneForward, "a prompt is one forward"); + RunTest(TestRestoreLoadsTheSnapshotWithoutReplaying, "restore loads the snapshot without replaying"); + RunTest(TestRewindingElsewhereReplaysInOnePass, "rewinding elsewhere replays in one pass"); + RunTest(TestRewindingOntoTheCheckpointLoadsItsSnapshot, "rewinding onto the checkpoint loads its snapshot"); + RunTest(TestRewindingOntoADivergedCheckpointReplays, "rewinding onto a diverged checkpoint replays"); + RunTest(TestAFailedRestorePoisonsTheEngine, "a failed restore poisons the engine"); + RunTest(TestAFailedForwardPoisonsTheEngineAndDropsTheCheckpoint, + "a failed forward poisons the engine and drops the checkpoint"); + RunTest(TestCheckpointingTheSameStateTwiceSavesOnce, "checkpointing the same state twice saves once"); + RunTest(TestRestoringRightAfterACheckpointLoadsNothing, "restoring right after a checkpoint loads nothing"); + RunTest(TestShrinkingBelowTheCheckpointDropsIt, "shrinking below the checkpoint drops it"); + RunTest(TestGrowingTheMaximumRecreatesTheModelAndKeepsTheSequence, "growing the maximum keeps the sequence"); + RunTest(TestShrinkingWithinTheCapacityKeepsTheModel, "shrinking within the capacity keeps the model"); + std::cout << "All Qwen3.6 rai tests passed\n"; + return 0; +} diff --git a/src/test/qwen3_rai/test_qwen3_rai.cpp b/src/test/qwen3_rai/test_qwen3_rai.cpp index 68f149d67..b99d041a0 100644 --- a/src/test/qwen3_rai/test_qwen3_rai.cpp +++ b/src/test/qwen3_rai/test_qwen3_rai.cpp @@ -236,6 +236,17 @@ std::vector Kinds(std::size_t from = 0) { return kinds; } +void TestStreamEnablesTheQwen3BundleBeforeAnyKernelQuery() { + Harness h; + const auto& state = fake_corelib::GetState(); + TEST_REQUIRE(state.kernel_bundles == std::vector({"qwen3"})); + const auto& log = state.call_log; + const auto enabled = std::find(log.begin(), log.end(), "ryzenai_corelib_stream_enable_kernel_bundle"); + const auto queried = std::find(log.begin(), log.end(), "ryzenai_corelib_matmul_bf16_enum_kernels"); + TEST_REQUIRE(queried != log.end()); + TEST_REQUIRE(enabled < queried); +} + void TestLoadPacksEveryWeightAtItsGroup() { Harness h; const auto& state = fake_corelib::GetState(); @@ -297,6 +308,18 @@ void TestDecodeRunsOneRowAtThePosition() { TEST_REQUIRE(h.engine->get_current_context_length() == 4); } +void TestTheVProjectionWritesTheCacheThroughAStridedWindow() { + Harness h; + (void)h.engine->forward(1); + const auto& windows = fake_corelib::GetState().tensor_windows; + const auto v = std::find_if(windows.begin(), windows.end(), + [](const auto& w) { return w.shape.size() == 3; }); + TEST_REQUIRE(v != windows.end()); + TEST_REQUIRE(v->strides.size() == 3 && v->strides[2] == 1); + TEST_REQUIRE(v->strides[1] == v->shape[2]); + TEST_REQUIRE(v->strides[0] == 4096 * v->strides[1]); +} + void TestContinuingPrefillReplaysTheConversationFromZero() { Harness h; std::vector first(10, 3); @@ -382,9 +405,13 @@ int main() { RunTest(TestUntiedPackageReadsOutputWeight, "untied package reads output.weight"); RunTest(TestContractMismatchesAreNamed, "contract mismatches are named"); RunTest(TestWrongArchitectureAndMissingTensorsAreRefused, "wrong architecture and missing tensors are refused"); + RunTest(TestStreamEnablesTheQwen3BundleBeforeAnyKernelQuery, + "stream enables the qwen3 bundle before any kernel query"); RunTest(TestLoadPacksEveryWeightAtItsGroup, "load packs every weight at its group"); RunTest(TestPrefillRunsEightDispatchesALayerAtTheBucket, "prefill runs eight dispatches a layer at the bucket"); RunTest(TestDecodeRunsOneRowAtThePosition, "decode runs one row at the position"); + RunTest(TestTheVProjectionWritesTheCacheThroughAStridedWindow, + "the V projection writes the cache through a strided window"); RunTest(TestContinuingPrefillReplaysTheConversationFromZero, "continuing prefill replays from zero"); RunTest(TestRestoreRewindsTheReplayHistory, "restore rewinds the replay history"); RunTest(TestUnknownHistoryCannotBeContinued, "unknown history cannot be continued"); From 1005b2edf28831520540f6f157afb35361ac0bf1 Mon Sep 17 00:00:00 2001 From: "Zhang, Chi" Date: Thu, 15 Oct 2026 21:14:53 -0700 Subject: [PATCH 02/10] refactor(rai): run Qwen3, Qwen3.5 and Gemma 4 as corelib's own models Qwen3, Qwen3.5 and Gemma 4 on the `rai` backend stop driving corelib op by op and run as corelib's own models (`ryzenai_corelib_model_*`), the way Qwen3.6 already does. All four share one engine and one backend; each model is a short spec. Phi-4 is unchanged (corelib's `phi4mini` takes ONNX only, and FastFlowLM ships a GGUF). - `flm::rai::CorelibModel` (`rai/corelib_model`): a `causal_lm` over a corelib model, generalized from the Qwen3.6 engine. `ModelSpec` states where a multi-row pass may start (`Continuation`: anywhere, only at 0, or at 0 / a 4096-row page boundary) and how to rewind (`Rewind`: seek, state snapshot, or reset and replay). Token ids and the context limit are checked before corelib is called; a failed forward, restore or reset poisons the engine. - `flm::rai::Factory(BackendSpec)` (`rai/corelib_model_backend`): refuses preemption and a context outside the model's rule, reads the package's BOS/EOS ids, and resolves the GGUF before the runtime or weights are touched. - Removed: the per-op engines, GGUF readers, shape plans and host math of Qwen3, Qwen3.5 and Gemma 4; `rai/gguf_file`, `rai/weight_cache`, the Q8/RMSNorm parts of `rai/host_ops`, and 14 unused corelib entry points. corelib packs the weights on every load; old `.rai-weights.bin/.json` files can be deleted. - Tests: new `src/test/rai_model` (engine and backend against the fake corelib; the Qwen3.6 tests move here). Gemma 4's registry, catalog and server-gate tests move to `test_gemma4_registry.cpp`; the op-level Gemma 4 and Qwen3 suites are gone. Gemma 4 is slower on corelib 0.26's built-in model than on the per-op engine (prefill ~1.5-2x, decode 24-29% slower, slower loads). The cause is in corelib's host per-layer-embedding projection and serial weight packing; FastFlowLM picks up the fix with a corelib version bump. Co-authored-by: Cursor --- docs/docs/models/qwen.md | 2 + src/common/models/README.md | 56 +- .../models/gemma4/rai/aie_next/gemma4_rai.cpp | 1929 --------- .../rai/aie_next/gemma4_rai_backend.cpp | 250 +- .../gemma4/rai/aie_next/gemma4_rai_gguf.cpp | 503 --- .../gemma4/rai/aie_next/gemma4_rai_host.cpp | 501 --- .../rai/aie_next/gemma4_rai_shape_plan.cpp | 374 -- .../models/qwen3/rai/aie_next/qwen3_rai.cpp | 667 --- .../qwen3/rai/aie_next/qwen3_rai_backend.cpp | 116 +- .../qwen3/rai/aie_next/qwen3_rai_config.cpp | 41 - .../qwen3/rai/aie_next/qwen3_rai_gguf.cpp | 182 - .../rai/aie_next/qwen3_rai_shape_plan.cpp | 83 - .../models/qwen35/rai/aie_next/qwen35_rai.cpp | 746 ---- .../rai/aie_next/qwen35_rai_backend.cpp | 81 +- .../qwen35/rai/aie_next/qwen35_rai_gguf.cpp | 355 -- .../models/qwen36/rai/aie_next/qwen36_rai.cpp | 289 -- .../rai/aie_next/qwen36_rai_backend.cpp | 106 +- src/common/rai/corelib_model.cpp | 352 ++ src/common/rai/corelib_model_backend.cpp | 205 + src/common/rai/gguf_file.cpp | 512 --- src/common/rai/host_ops.cpp | 91 - src/common/rai/rai_sources.cmake | 4 +- src/common/rai/weight_cache.cpp | 149 - .../models/gemma4/rai/aie_next/gemma4_rai.hpp | 126 - .../rai/aie_next/gemma4_rai_backend.hpp | 50 +- .../rai/aie_next/gemma4_rai_constants.hpp | 75 - .../gemma4/rai/aie_next/gemma4_rai_gguf.hpp | 283 -- .../gemma4/rai/aie_next/gemma4_rai_host.hpp | 160 - .../rai/aie_next/gemma4_rai_shape_plan.hpp | 153 - .../models/qwen3/rai/aie_next/qwen3_rai.hpp | 48 - .../qwen3/rai/aie_next/qwen3_rai_backend.hpp | 3 - .../qwen3/rai/aie_next/qwen3_rai_config.hpp | 107 - .../qwen3/rai/aie_next/qwen3_rai_gguf.hpp | 60 - .../rai/aie_next/qwen3_rai_shape_plan.hpp | 36 - .../models/qwen35/rai/aie_next/qwen35_rai.hpp | 43 - .../rai/aie_next/qwen35_rai_backend.hpp | 1 - .../qwen35/rai/aie_next/qwen35_rai_gguf.hpp | 103 - .../models/qwen36/rai/aie_next/qwen36_rai.hpp | 46 - .../rai/aie_next/qwen36_rai_backend.hpp | 1 - src/include/rai/corelib_api.hpp | 14 - src/include/rai/corelib_model.hpp | 87 + src/include/rai/corelib_model_backend.hpp | 36 + src/include/rai/corelib_object.hpp | 6 - src/include/rai/gguf_file.hpp | 101 - src/include/rai/host_ops.hpp | 11 - src/include/rai/weight_cache.hpp | 80 - src/test/gemma4_rai/CMakeLists.txt | 438 +- src/test/gemma4_rai/fake_corelib.cpp | 1252 ------ src/test/gemma4_rai/fake_corelib.hpp | 920 ---- src/test/gemma4_rai/gemma4_gguf_fixture.hpp | 884 ---- src/test/gemma4_rai/test_gemma4_backend.cpp | 1523 ------- src/test/gemma4_rai/test_gemma4_config.cpp | 530 --- .../gemma4_rai/test_gemma4_derivations.cpp | 258 -- src/test/gemma4_rai/test_gemma4_engine.cpp | 3844 ----------------- src/test/gemma4_rai/test_gemma4_gguf.cpp | 361 -- src/test/gemma4_rai/test_gemma4_host.cpp | 501 --- src/test/gemma4_rai/test_gemma4_real_gguf.cpp | 632 --- src/test/gemma4_rai/test_gemma4_registry.cpp | 615 +++ .../gemma4_rai/test_gemma4_shape_plan.cpp | 889 ---- src/test/gemma4_rai/test_gemma4_tokenizer.cpp | 15 +- src/test/phi4_rai/fake_corelib.cpp | 43 +- src/test/phi4_rai/fake_corelib.hpp | 3 + src/test/qwen36_rai/CMakeLists.txt | 42 - src/test/qwen36_rai/test_qwen36_rai.cpp | 236 - src/test/qwen3_rai/CMakeLists.txt | 49 - src/test/qwen3_rai/test_qwen3_rai.cpp | 423 -- src/test/rai_model/CMakeLists.txt | 49 + src/test/rai_model/test_corelib_model.cpp | 482 +++ .../rai_model/test_corelib_model_backend.cpp | 261 ++ 69 files changed, 2288 insertions(+), 21186 deletions(-) delete mode 100644 src/common/models/gemma4/rai/aie_next/gemma4_rai.cpp delete mode 100644 src/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp delete mode 100644 src/common/models/gemma4/rai/aie_next/gemma4_rai_host.cpp delete mode 100644 src/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp delete mode 100644 src/common/models/qwen3/rai/aie_next/qwen3_rai.cpp delete mode 100644 src/common/models/qwen3/rai/aie_next/qwen3_rai_config.cpp delete mode 100644 src/common/models/qwen3/rai/aie_next/qwen3_rai_gguf.cpp delete mode 100644 src/common/models/qwen3/rai/aie_next/qwen3_rai_shape_plan.cpp delete mode 100644 src/common/models/qwen35/rai/aie_next/qwen35_rai.cpp delete mode 100644 src/common/models/qwen35/rai/aie_next/qwen35_rai_gguf.cpp delete mode 100644 src/common/models/qwen36/rai/aie_next/qwen36_rai.cpp create mode 100644 src/common/rai/corelib_model.cpp create mode 100644 src/common/rai/corelib_model_backend.cpp delete mode 100644 src/common/rai/gguf_file.cpp delete mode 100644 src/common/rai/weight_cache.cpp delete mode 100644 src/include/models/gemma4/rai/aie_next/gemma4_rai.hpp delete mode 100644 src/include/models/gemma4/rai/aie_next/gemma4_rai_constants.hpp delete mode 100644 src/include/models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp delete mode 100644 src/include/models/gemma4/rai/aie_next/gemma4_rai_host.hpp delete mode 100644 src/include/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.hpp delete mode 100644 src/include/models/qwen3/rai/aie_next/qwen3_rai.hpp delete mode 100644 src/include/models/qwen3/rai/aie_next/qwen3_rai_config.hpp delete mode 100644 src/include/models/qwen3/rai/aie_next/qwen3_rai_gguf.hpp delete mode 100644 src/include/models/qwen3/rai/aie_next/qwen3_rai_shape_plan.hpp delete mode 100644 src/include/models/qwen35/rai/aie_next/qwen35_rai.hpp delete mode 100644 src/include/models/qwen35/rai/aie_next/qwen35_rai_gguf.hpp delete mode 100644 src/include/models/qwen36/rai/aie_next/qwen36_rai.hpp create mode 100644 src/include/rai/corelib_model.hpp create mode 100644 src/include/rai/corelib_model_backend.hpp delete mode 100644 src/include/rai/gguf_file.hpp delete mode 100644 src/include/rai/weight_cache.hpp delete mode 100644 src/test/gemma4_rai/fake_corelib.cpp delete mode 100644 src/test/gemma4_rai/fake_corelib.hpp delete mode 100644 src/test/gemma4_rai/gemma4_gguf_fixture.hpp delete mode 100644 src/test/gemma4_rai/test_gemma4_backend.cpp delete mode 100644 src/test/gemma4_rai/test_gemma4_config.cpp delete mode 100644 src/test/gemma4_rai/test_gemma4_derivations.cpp delete mode 100644 src/test/gemma4_rai/test_gemma4_engine.cpp delete mode 100644 src/test/gemma4_rai/test_gemma4_gguf.cpp delete mode 100644 src/test/gemma4_rai/test_gemma4_host.cpp delete mode 100644 src/test/gemma4_rai/test_gemma4_real_gguf.cpp create mode 100644 src/test/gemma4_rai/test_gemma4_registry.cpp delete mode 100644 src/test/gemma4_rai/test_gemma4_shape_plan.cpp delete mode 100644 src/test/qwen36_rai/CMakeLists.txt delete mode 100644 src/test/qwen36_rai/test_qwen36_rai.cpp delete mode 100644 src/test/qwen3_rai/CMakeLists.txt delete mode 100644 src/test/qwen3_rai/test_qwen3_rai.cpp create mode 100644 src/test/rai_model/CMakeLists.txt create mode 100644 src/test/rai_model/test_corelib_model.cpp create mode 100644 src/test/rai_model/test_corelib_model_backend.cpp diff --git a/docs/docs/models/qwen.md b/docs/docs/models/qwen.md index df513a24e..03d4bc306 100644 --- a/docs/docs/models/qwen.md +++ b/docs/docs/models/qwen.md @@ -117,6 +117,8 @@ The same corelib path also runs `gemma4-it-rai:e2b` and `gemma4-it-rai:e4b` (tex - **Checkpoints take host memory.** A checkpoint (what lets a conversation rewind without replaying) copies the KV cache into host memory, in whole 4096-token pages up to the current position: about 80 MiB per filled page, so up to about 5 GiB at a full 262144-token context. - **`gemma4-it-rai:e2b` / `gemma4-it-rai:e4b` — up to 16384 tokens**, the length corelib has validated paged Gemma 4 to, and never past the GGUF's own `gemma4.context_length`. The default is 16384. Up to 4096 any length works; above 4096 the context must be a multiple of 4096. The KV cache for the whole context is allocated up front, when the model loads. +Qwen3, Qwen3.5, Qwen3.6 and Gemma 4 run as corelib's own models, which pack their weights on every load: there is no packed-weight cache beside the GGUF. A `.rai-weights.bin` / `.json` pair left there by an earlier build is no longer read and can be deleted. + ```powershell flm pull qwen3-rai:4b flm run qwen3-rai:4b diff --git a/src/common/models/README.md b/src/common/models/README.md index 9d5d1922c..dddf1723f 100644 --- a/src/common/models/README.md +++ b/src/common/models/README.md @@ -1,9 +1,16 @@ # Adding a model on the rai backend How to bring up a new model family on **rai**, the backend that reaches its -kernels through ryzenai-corelib, using Phi-4 as the worked example. Qwen3, Gemma 4 -and Qwen3.5 follow the same layout, but every file named below has a `phi4` -counterpart you can read straight through. +kernels through ryzenai-corelib. There are two ways, and the first is the one +to reach for: + +- **corelib runs the model whole** (it ships `ryzenai_corelib_model_`): + write one `BackendSpec` (its `.model` is the `ModelSpec`) and hand it to + `flm::rai::Factory` — see §2a. Qwen3, Qwen3.5, Qwen3.6 and Gemma 4 are built + this way, each in one short + `_rai_backend.cpp`. +- **corelib has only the ops**: write the engine yourself, op by op, as + Phi-4 does — the rest of this document, from §2 on, walks through it. A backend names *where the kernels come from*: `flm` is FastFlowLM's own kernel flow, `rai` is corelib. That is a separate axis from which silicon the host has, @@ -35,7 +42,7 @@ This is a contributor document. For *using* a backend once it exists — `--back |---|---| | **Platform** | Windows and Linux. Both configure and build; only Windows has been run on hardware, and the `src/test/phi4_rai` suite is still Windows-only. | | **Hardware** | An aie_next NPU. There is no simulator; a wrong shape shows up as garbage output, not an error. | -| **corelib headers** | The corelib ABI is still pre-1.0, so FastFlowLM requires an **exact `0.26.0`** match on major, minor and patch. [`corelib_api.hpp`](../../include/rai/corelib_api.hpp) `#error`s on any other version — deliberately, because the C ABI has changed shape between patch releases. The version is queried before any other entry point, so a mismatched runtime reports a version error rather than a missing symbol. corelib's default build gives one `ryzenai_corelib_model_` library per model, carrying its kernels, which must sit beside corelib. Each engine enables its model's kernel bundle on the stream it opens. | +| **corelib headers** | The corelib ABI is still pre-1.0, so FastFlowLM requires an **exact `0.26.0`** match on major, minor and patch. [`corelib_api.hpp`](../../include/rai/corelib_api.hpp) `#error`s on any other version — deliberately, because the C ABI has changed shape between patch releases. The version is queried before any other entry point, so a mismatched runtime reports a version error rather than a missing symbol. corelib's default build gives one `ryzenai_corelib_model_` library per model, carrying its kernels, which must sit beside corelib. A per-op engine (Phi-4) enables its model's kernel bundle on the stream it opens; a model corelib runs whole enables its own. | | **Weights** | A GGUF the vendor kernels can requantize. Phi-4 uses Q8_0; each weight is a `*_weights_create` call whose qweight source states `gguf_q8_0`. | Configure with: @@ -78,6 +85,47 @@ entries, not in the backend id. --- +## 2a. A model corelib runs whole + +[`rai/corelib_model.hpp`](../../include/rai/corelib_model.hpp) is a `causal_lm` +over `ryzenai_corelib_model_*`; [`rai/corelib_model_backend.hpp`](../../include/rai/corelib_model_backend.hpp) +is the `ModelBackend` around it. A family needs only its spec: + +```cpp +flm::backend::BackendFactory rai_factory() { + return rai::Factory({ + .model = {"qwen3", "Qwen3", kRaiContextLimit, rai::Continuation::kFromStart, + rai::Rewind::kSeek}, + .context_rule = "1..4096", + .eos_ids = {151645, 151643}, + }); +} +``` + +Two fields say how the model behaves, and both come from what corelib's driver +accepts (read its `models/.cpp`): + +| field | value | when | +|---|---|---| +| `continuation` | `kAnywhere` | the driver takes `rows > 1` at any position | +| | `kFromStart` | only at position 0; a later prompt re-runs the conversation | +| | `kFromPage` | at 0 or on a 4096-row page boundary of a paged cache | +| `rewind` | `kSeek` | attention only: each pass states its position | +| | `kSnapshot` | recurrent state in state slots: `model_state_save/restore` | +| | `kReplay` | recurrent state outside any slot: reset and replay | + +Leave `eos_ids` empty to read the stop tokens from the package, and set +`read_bos` when `tokenizer_config.json` names a BOS but no id states it. +corelib packs the weights on every load; there is no packed-weight cache on +this path. + +The header (`_rai_backend.hpp`) keeps the `inline rai_traits()` and the +context constants; registration in `builtin_backends.cpp` is unchanged. +Tests: [`test/rai_model/`](../../test/rai_model/) covers the engine and the +backend against the fake corelib; a new family adds no engine test of its own. + +--- + ## 2. Lay out the files ``` diff --git a/src/common/models/gemma4/rai/aie_next/gemma4_rai.cpp b/src/common/models/gemma4/rai/aie_next/gemma4_rai.cpp deleted file mode 100644 index 8b1671f21..000000000 --- a/src/common/models/gemma4/rai/aie_next/gemma4_rai.cpp +++ /dev/null @@ -1,1929 +0,0 @@ -#include "models/gemma4/rai/aie_next/gemma4_rai.hpp" - -#include "models/gemma4/rai/aie_next/gemma4_rai_constants.hpp" -#include "models/gemma4/rai/aie_next/gemma4_rai_host.hpp" -#include "models/gemma4/rai/aie_next/gemma4_rai_shape_plan.hpp" -#include "rai/corelib_object.hpp" -#include "rai/gguf_file.hpp" -#include "rai/weight_cache.hpp" -#include "rai/weight_source.hpp" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace flm::gemma4 { -namespace { -using namespace flm::corelib; -using namespace flm::rai; - -std::string Blk(std::int64_t layer, const char* suffix) { - return "blk." + std::to_string(layer) + suffix; -} - -/// \brief `f32` as the BF16 the kernels read -std::vector ToBf16(std::span values) { - std::vector out(values.size()); - FloatsToBf16(values, out); - return out; -} - -/// \brief `source`, a [rows, cols] row-major matrix, as a [cols, rows] one -/// \note THE ONE PLACE A GGUF MAPPING IS RE-LAID-OUT IN THIS ENGINE. Every -/// other corelib port takes the file's own memory order (see the -/// layout table in the review of C7/R1: matmul's `blocks`, ssmlp's -/// gate/up/down, both norm ports are all correct as mapped). `ple` is -/// the exception, and the transpose is materialised rather than -/// described because `ryzenai_corelib_ple_bf16_weights_create` takes a -/// dense FP32 source with no stride. -std::vector Transposed(std::span source, std::int64_t rows, - std::int64_t cols) { - const auto row_count = static_cast(rows); - const auto column_count = static_cast(cols); - std::vector out(row_count * column_count); - for (std::size_t r = 0; r < row_count; ++r) { - const auto* in_row = source.data() + r * column_count; - for (std::size_t c = 0; c < column_count; ++c) - out[c * row_count + r] = in_row[c]; - } - return out; -} - -/// \brief a gamma of exactly 1.0, `k` wide -/// \note V's norm. The file carries NO learned weight for it -/// (llama.cpp src/models/gemma4.cpp:256), and `rmsnorm_bf16` ships no -/// weightless variant, so "normalize only" is spelled as a vector of -/// ones -- which is exact in BF16. The reference driver keeps exactly -/// this vector, for exactly this reason (`gemma4_driver.py` -/// `_ones_norm_weights`): "the cheap wrong implementation omits the -/// step entirely and still produces fluent text; on this model that is -/// a whole 1/sqrt(mean(v^2)) missing from every attended value." -std::vector OnesBf16(std::int64_t k) { - const std::vector ones(static_cast(k), 1.0f); - return ToBf16(ones); -} - -/// \brief everything besides the GGUF and corelib version that shapes the -/// packed bytes -std::string CacheLayout(const Gemma4Config& cfg) { - return "gemma4 " + std::string(cfg.row_name) + " group=" + std::to_string(cfg.group) + - " head_group=" + std::to_string(cfg.head_group) + - " ple_group=" + std::to_string(cfg.ple_group) + - " layout_revision=" + std::to_string(kPackedLayoutRevision); -} - -enum class SlotKind { Matmul, SsMlp, RmsNorm, Ple }; - -/// \brief one packed weight: how to make it, where to put it, how to cache it -/// -/// ONE TABLE, WALKED THE SAME WAY BOTH DIRECTIONS. The pack path runs -/// `pack()` for every slot across a thread pool; the cache path calls the -/// matching `..._weights_create` with the cached blob as its one `packed` -/// source and `assign()`s the handle. Because both walk the -/// same vector in the same order, a cache index entry always refers to the -/// weight it was written from -- which is the property Phi-4's engine gets -/// from a hand-written `slot % 5` scheme and Gemma 4 cannot, because its -/// per-layer slot count is not constant (ten on an owning layer, six on a -/// sharing one). -/// -/// ASSIGNMENT IS BY SLOT, NEVER BY COMPLETION ORDER. `pack()` writes into -/// the one member it owns; the pool's workers finish in whatever order they -/// finish in and nothing appends to a shared list. -struct WeightSlot { - SlotKind kind{}; - std::string label; - /// \brief the layer this slot belongs to, or -1 for a model-level one - /// \note Only ple uses it, to find its blob. Kept on every slot so the - /// labels and the blob index cannot drift apart. - std::int64_t layer{-1}; - - ryzenai_corelib_matmul_bf16_weights_desc matmul{}; - ryzenai_corelib_ssmlp_bf16_weights_desc ssmlp{}; - ryzenai_corelib_rmsnorm_bf16_weights_desc rmsnorm{}; - ryzenai_corelib_ple_bf16_weights_desc ple{}; - - /// \brief pack this weight from the GGUF and store it - std::function pack; - /// \brief adopt an already-created handle (the cache path) - std::function assign; - /// \brief the live handle, for the cache write; null before packing - std::function handle; -}; - -} // namespace - -struct gemma4_rai::Impl { - /// DECLARATION ORDER IS DESTRUCTION ORDER, REVERSED, AND IT IS LOAD- - /// BEARING. Members are destroyed bottom-up, so the packed weights go - /// first, then the stream, then the API and the package, and the runtime - /// last of all. corelib objects must not outlive the API they came from, - /// and A PACKED OBJECT MUST BE RELEASED BEFORE THE STREAM IT WAS PLANNED - /// AGAINST. - /// - /// THAT LAST RULE IS MEASURED, NOT TIDINESS. corelib access-violates - /// inside `ryzenai_corelib_object_release` when an object outlives its - /// Stream: with two models in one process and the first torn down - /// stream-first, the reference driver printed `Windows fatal exception: - /// access violation` in 7 of 10 runs, up to 221 lines in one, against - /// 0 of 10 for the packed-objects-first order. Box 43, - /// `test_gemma4_weights.py`, twenty runs over the two orders, all of it - /// in `t17r1d.log`. - /// - /// WHAT THE EXIT CODES SAY, STATED AS MEASURED. An earlier version of - /// this comment said "AND THE PROCESS EXITED 0 EVERY TIME", and the - /// source does not support that. What the corelib model notes record - /// is that across ALL TWENTY runs -- both arms - /// together -- nineteen sessions exited 0 with every test PASSED and THE - /// TWENTIETH EXITED 1, with no violations at all. The document does not - /// break exit codes down by arm, so nothing establishes what the - /// stream-first ten did. "Nothing raised to the caller" is narrower - /// still: what was measured is that a recorder in `Object.__del__` - /// caught ZERO of the 833 violations those runs printed, and the source - /// says in terms that a mechanism must not be re-derived from that zero. - /// Every figure here is one shared box that takes no lock on the NPU, - /// which is why the driver's own docstring declines to repeat any of - /// them. - /// - /// THE CONCLUSION IS UNCHANGED, WHICH IS WHY THE MEASUREMENT IS WORTH - /// GETTING RIGHT. At 19 of 20 an exit-code test asserts nothing either - /// way, so the guard here has to be structural. Their driver met this - /// only because a second model entered one process; FLM loads and - /// unloads models routinely, so it is a shipping path here. The - /// underlying defect is corelib's, is undiagnosed and unfixed, and is - /// not touched here -- this engine simply does not reach it. - /// - /// THERE IS DELIBERATELY NO `~Impl()`, and that is the C9 decision rather - /// than an omission. A destructor body runs BEFORE member destruction, so - /// a hand-written `weights.clear(); ...; stream.reset();` list would - /// release the stream before any member the list FORGOT -- which is - /// precisely the fault, arriving silently the first time somebody adds a - /// member and does not update the list. Declaration order covers every - /// member automatically and has exactly one failure mode: a new - /// corelib-owning member declared above `stream`. That one failure mode - /// is what `TestEveryCorelibOwningMemberIsDeclaredAfterTheStream` in - /// src/test/gemma4_rai/test_gemma4_engine.cpp reads this block to catch. - /// A behavioural test cannot: the fault is intermittent and raises - /// nothing, so one that waited for a crash would pass on a good day with - /// bad code. - /// - /// IF YOU ADD A MEMBER: a corelib handle, or anything holding one, goes - /// BELOW `stream`. Host memory corelib may have BOUND goes ABOVE the - /// weights that bound it, and needs its own line in that test, because - /// nothing in the type says which buffers those are. - std::shared_ptr runtime; - std::shared_ptr package; - std::shared_ptr api; - Gemma4Config shape; - std::uint32_t max_length; - /// \brief context longer than one flat_mha window: paged KV, PDI (8, 33) - bool paged{}; - /// \brief rows the KV cache and rotary tables hold - /// \note Unpaged this stays kMaxSequenceLength even when the caller asked - /// for less, which is the pitch flat_mha addresses. Paged it is the - /// constructed context, a whole number of pages. - std::uint32_t cache_rows{}; - int position{}; - std::optional saved; - bool poisoned{}; - bool from_cache{}; - std::size_t slot_count{}; - - UniqueStream stream; - Gemma4ShapePlan plan; - - struct LayerWeights { - /// Present on every layer. - UniqueMatMulWeights q, o; - UniqueRmsNormWeights q_norm, post_attn_norm; - UniqueSsMlpWeights mlp; - UniquePleWeights ple; - /// Present ONLY on a layer that owns its KV cache. A sharing layer - /// never projects or norms K and V -- the kvshare attention kernel - /// "is not handed K at all" -- so these stay empty, and that is the - /// config's decision, not the file's: the tensors exist on every - /// layer of every shipped conversion. - UniqueMatMulWeights k, v; - UniqueRmsNormWeights k_norm, v_norm; - }; - std::vector layers; - - /// \brief every activation buffer, allocated once at the WIDER geometry - /// - /// SIZED AT THE WIDEST HEAD THIS MODEL HAS, THEN WINDOWED PER LAYER ON THE - /// COLUMN AXIS. Gemma 4 runs two geometries in one stack -- on E2B Q, O - /// and the attention output are 4096 wide on a full layer and 2048 on a - /// sliding one, K/V 512 and 256 -- and each op reads its row and column - /// count off its input's shape, so a window is what says which kind of - /// layer this is. One buffer at the maximum costs a few MB more than the - /// largest single layer needs and removes any question of which kind wrote - /// a buffer last. - /// - /// THAT IS THE COLUMN AXIS AND IT DOES NOT LICENSE THE ROW AXIS. A padded - /// buffer is not a window onto a bigger one: every op here reads and - /// writes its whole padded extent, so the row count a window declares is - /// the row count the kernel runs, not a bound on it. - /// - /// FOUR HIDDEN-WIDTH BUFFERS AND NO SWAP, which is where this differs from - /// Phi-4. Phi-4 ping-pongs `hidden`/`normalized` because its fused MLP - /// writes the next layer's input into the buffer it just read; here `ple` - /// produces BOTH the next residual and the next normed input, so four - /// roles rotate through four fixed buffers and every op in a layer has - /// four distinct operands: - /// - /// hidden the normed input, then o_proj's output, then the plane - /// ssmlp does NOT write, then ple's norm_out - /// skip_sum post_attention_norm's output, ssmlp's input - /// ple_out ssmlp's block output (plane 0), which is ple's `x` - /// residual the residual stream, which ple overwrites with the next - UniqueTensor hidden, residual, skip_sum, ple_out; - /// \brief the projections' outputs, and QK-Norm's separate destinations - /// \note `q_normed`/`k_normed` are NOT a buffer that could be saved by - /// normalizing in place. corelib.h measures in-place `rmsnorm_bf16` - /// CORRUPTING at M = 1024 and M = 2048 at k=128 -- on EXACT-kernel - /// dispatches, so "a kernel ships at this shape" is not the test -- - /// with a failure shape ("a handful of consecutive rows out of a - /// thousand") that keeps correlation high and rarely moves an - /// argmax. A prefill Q-norm runs at `padded * q_heads` rows, which - /// is exactly that range. - /// \note `v_buf` is the buffer Phi-4 does not need. Gemma 4 RMS-normalizes - /// V between the projection and the cache, and no norm can scatter, - /// so the projection lands here and THE NORM WRITES THE CACHE. - UniqueTensor q_buf, k_buf, v_buf, q_normed, k_normed, attn_out; - /// \brief lm_head's input and output, one row each - /// \note Not a simplification: lm_head is `hidden -> vocab` and ships at - /// M == 1 alone at this model's (8, 17) PDI pair. - UniqueTensor last_hidden, logits; - /// \brief the per-layer embedding input planes, one tensor per layer - /// \note One per layer rather than one [layers, rows, ple_dim] tensor so - /// that a layer's slice is a single write and a single bind. - std::vector ple_in; - - /// \brief the KV caches: `kv_layers` of them, NOT one per layer, and NOT - /// uniform in width - /// - /// CACHE `i` IS AS WIDE AS LAYER `i`'S HEAD. On E2B caches 4, 9 and 14 are - /// 512 and the other twelve are 256; on E4B it is 5, 11, 17 and 23. Sizing - /// them all at one head width "either wastes half the allocation or, - /// worse, hands the full-attention layers a cache with the wrong row pitch - /// -- which the kernel would address happily". - /// - /// PINNED TO `kMaxSequenceLength`, NEVER GROWN TO FIT, and that is - /// corelib's rule rather than a convenience: "the kernel addresses them at - /// that row pitch, so a cache sized to the current sequence changes stride - /// every token and reads the wrong rows". The pitch it means is - /// `flat_mha_bf16_desc::max_seq`, which this engine's shape plan sets to - /// `kMaxSequenceLength` -- so these two numbers are one number, and - /// allocating at the constructor's `max_length` instead would desynchronize - /// them. `max_length` bounds what a caller may REQUEST, not what is - /// allocated. - std::vector k_cache, v_cache; - - /// \brief the two rotary table pairs, sliding then full-attention, as host - /// views over engine-owned storage - /// \note Two pairs rather than one rewritten per layer: corelib memoizes - /// the derived table on the view's identity. Written once at load and - /// never touched again. corelib copies nothing, so the storage is - /// declared first and outlives the views. - std::vector rope_cos_sliding, rope_sin_sliding, rope_cos_full, rope_sin_full; - /// \brief the tokens behind `position`, so a multi-row continuation can be - /// re-run from position 0 (attention ships no chunked prefill) - std::vector history; - UniqueHostView cos_sliding, sin_sliding, cos_full, sin_full; - - /// \brief the mapped tensors a forward pass reads on the HOST - /// \note Views into the GGUF mapping, which `package` keeps alive for the - /// engine's life. `per_layer_token_embd` is 2.35 billion Q8_0 - /// parameters on E2B -- materializing it is not an option at any - /// point, so it stays mapped and is gathered per token. - GgufTensor embedding_table, ple_table; - GgufBf16Tensor ple_projection; - GgufFloatTensor ple_projection_norm; - - /// \brief layer 0's own attention norm, the model's ONE standalone - /// RMSNorm - /// \note Every other input norm is the previous layer's `ple` second - /// output (`norm_out = RMSNorm(out, next_norm)`), which crosses the - /// block boundary -- so nothing between two layers applies one, and - /// layer 0 has no block before it to produce its own. - UniqueRmsNormWeights input_norm; - /// \brief lm_head, TIED to token_embd and the ONE weight at head_group - UniqueMatMulWeights lm_head; - - Impl(std::shared_ptr package_in, - std::shared_ptr runtime_in, const LM_Config&, - std::optional shape_override, std::uint32_t maximum) - : runtime(std::move(runtime_in)), package(std::move(package_in)), - api(runtime ? runtime->api() : nullptr), max_length(maximum) { - if (!package) throw std::invalid_argument("Gemma 4 GGUF package is null"); - if (!runtime || !api) throw std::invalid_argument("corelib runtime is null"); - if (!ContextFits(maximum)) { - throw std::invalid_argument( - "Gemma 4 maximum length must be in 1..4096, or a multiple of 4096 up to " + - std::to_string(kMaxPagedContext)); - } - paged = maximum > static_cast(kPageSize); - cache_rows = paged ? maximum : static_cast(kMaxSequenceLength); - shape = shape_override ? std::move(*shape_override) : package->Config(); - if (paged && static_cast(maximum) > shape.context_length) { - throw std::invalid_argument( - "Gemma 4 maximum length " + std::to_string(maximum) + - " is past this GGUF's context of " + std::to_string(shape.context_length)); - } - - const auto& cfg = shape; - const auto& file = package->File(); - const auto layer_count = static_cast(cfg.layers); - layers.resize(layer_count); - - // ---- host-side material, all of it resolved before any thread runs - // - // Every span below is either a view into the mapped GGUF (which the - // package keeps alive) or a vector that lives until the pool joins. - // Resolving them up front is what makes the creates independent of - // each other: a worker touches its own slot and nothing else. - const std::array epsilon_value{cfg.eps}; - const auto epsilon = ToBf16(epsilon_value); - - // ORIENTATION, AND IT IS [OUT, IN]. `GgufFile` reports a tensor's - // shape REVERSED from the GGUF's own dims order, and the file stores - // dims fastest-varying (input) first -- so `ShapeOf` answers - // [output, input] and every expected shape written here is in that - // convention. `blk.0.attn_q.weight` reads [1536, 2048] in the file - // and is demanded as {2048, 1536} below. It is the same convention - // phi4_rai.cpp uses for `token_embd.weight` ({vocab, hidden}), which - // is the comparison that settles it: Phi-4 has run against a real - // GGUF on hardware. - // - // These were all written the other way round until Task R1, because - // gemma4_gguf_fixture.hpp was transposed and this file was checked - // against the fixture. Every one of them is now asserted against the - // real gemma-4-E{2,4}B-it-Q8_0.gguf by test_gemma4_real_gguf.cpp. - // If one is ever wrong again, RequireQ8/RequireF32 throw naming the - // tensor and both shapes; nothing here can pick up a - // differently-shaped tensor and pack it quietly. - // - // The two parameters are named `out` and `in` rather than - // `rows`/`cols` deliberately: "rows" is the axis whose meaning the - // whole defect turned on. - const auto q8 = [&](const std::string& name, std::int64_t out, - std::int64_t in) { - const std::array want{out, in}; - return file.RequireQ8(name, want); - }; - const auto f32_1d = [&](const std::string& name, std::int64_t k) { - const std::array want{k}; - return file.RequireF32(name, want); - }; - const auto f32_2d = [&](const std::string& name, std::int64_t out, - std::int64_t in) { - const std::array want{out, in}; - return file.RequireF32(name, want); - }; - - const auto embedding = q8("token_embd.weight", cfg.vocab, cfg.hidden); - - // THE THREE MAPPED TENSORS A FORWARD PASS READS ON THE HOST, resolved - // here so a malformed file is named at load rather than at the first - // prompt. `per_layer_model_proj.weight` is the ONE BF16 tensor in a - // Gemma 4 conversion; the other two are Q8_0 and F32. - embedding_table = embedding; - ple_table = q8("per_layer_token_embd.weight", cfg.vocab, - cfg.layers * cfg.ple_dim); - { - const std::array want{cfg.layers * cfg.ple_dim, cfg.hidden}; - ple_projection = file.RequireBf16("per_layer_model_proj.weight", want); - } - ple_projection_norm = f32_1d("per_layer_proj_norm.weight", cfg.ple_dim); - - // `rope_freqs.weight`, ONE ENTRY PER ROTARY PAIR OF THE FULL HEAD. - // [256] on both shipped rows, which is `global_head_dim / 2` and NOT a - // 256-wide rotary -- `config.json`'s `partial_rotary_factor: 0.25` is - // a per-pair frequency SCALING, not a shorter rotary, and - // `gemma4.rope.dimension_count` (512) / `..._swa` (256) confirm full - // rotary on both geometries. Demanded at that exact length here, so a - // conversion that ships another one is named at load by RequireF32 - // rather than read past its end; MakeRopeTables checks the same thing - // again against the head it is building for. The driver checks it too. - const auto rope_factors = - f32_1d("rope_freqs.weight", cfg.global_head_dim / 2); - - // BF16 copies of every gamma that goes into a packer. Held in vectors - // rather than passed as mapped floats because both packers take BF16 - // and the file stores F32. - std::vector> ffn_norm(layer_count), - post_ffw_norm(layer_count), post_attn_norm(layer_count), - q_norm(layer_count), k_norm(layer_count); - auto input_norm_bf16 = ToBf16(f32_1d("blk.0.attn_norm.weight", cfg.hidden).values); - // One per DISTINCT head size, not one per layer: a sliding layer's - // head is 256 wide and a full-attention layer's 512, and `k` selects - // the rmsnorm kernel. - auto ones_sliding = OnesBf16(cfg.head_dim); - auto ones_full = OnesBf16(cfg.global_head_dim); - - std::vector q_weight(layer_count), k_weight(layer_count), - v_weight(layer_count), o_weight(layer_count), gate_weight(layer_count), - up_weight(layer_count), down_weight(layer_count); - std::vector inp_gate(layer_count), proj(layer_count), - ple_post_norm(layer_count), ple_next_norm(layer_count); - std::vector layer_scale(layer_count); - - for (std::size_t index = 0; index < layer_count; ++index) { - const auto layer = static_cast(index); - const bool owns = layer < cfg.kv_layers; - const auto head_dim = cfg.LayerHeadDim(layer); - const auto q_dim = cfg.q_heads * head_dim; - const auto kv_dim = cfg.kv_heads * head_dim; - const auto intermediate = cfg.layer_intermediate[index]; - - q_weight[index] = q8(Blk(layer, ".attn_q.weight"), q_dim, cfg.hidden); - // THE TWO PROJECTIONS WHOSE INPUT IS NOT `hidden` -- attn_output - // takes q_dim in, ffn_down takes `intermediate`. - o_weight[index] = q8(Blk(layer, ".attn_output.weight"), cfg.hidden, q_dim); - gate_weight[index] = q8(Blk(layer, ".ffn_gate.weight"), intermediate, cfg.hidden); - up_weight[index] = q8(Blk(layer, ".ffn_up.weight"), intermediate, cfg.hidden); - down_weight[index] = q8(Blk(layer, ".ffn_down.weight"), cfg.hidden, intermediate); - if (owns) { - k_weight[index] = q8(Blk(layer, ".attn_k.weight"), kv_dim, cfg.hidden); - v_weight[index] = q8(Blk(layer, ".attn_v.weight"), kv_dim, cfg.hidden); - k_norm[index] = ToBf16(f32_1d(Blk(layer, ".attn_k_norm.weight"), head_dim).values); - } - q_norm[index] = ToBf16(f32_1d(Blk(layer, ".attn_q_norm.weight"), head_dim).values); - post_attn_norm[index] = - ToBf16(f32_1d(Blk(layer, ".post_attention_norm.weight"), cfg.hidden).values); - // ssmlp's TWO NORMS ARE BOTH THIS LAYER'S. With - // `activation = 1, post_feedforward_layernorm = 1` the descriptor - // keys the `ssmlp_gemma_fusion_no_rms1_*` family, whose topology - // is NOT the silu one: the trailing norm is dropped and `norm1` - // becomes the post-feedforward gamma applied INSIDE the residual. - // So norm0 is `ffn_norm` and norm1 is `post_ffw_norm` and NEITHER - // belongs to another layer -- unlike the silu family, which is - // what Phi-4 is. Carrying Phi-4's assignment over here is silent. - // (gemma4_driver.py:1441-1449.) - ffn_norm[index] = ToBf16(f32_1d(Blk(layer, ".ffn_norm.weight"), cfg.hidden).values); - post_ffw_norm[index] = - ToBf16(f32_1d(Blk(layer, ".post_ffw_norm.weight"), cfg.hidden).values); - - // ple's two matrices, as mapped floats. - // - // TWO DIFFERENT CONVENTIONS MEET HERE AND IT IS WORTH BEING - // EXPLICIT. corelib takes `gate` as FP32 [k, n] row-major and - // `proj` as FP32 [n, k] row-major -- "BOTH MATRICES ARE GIVEN IN - // [IN, OUT] ORDER, which is the orientation a model stores and - // the transpose of what the kernel wants; the transpose happens - // here" (corelib.h, ple section). `GgufFile::ShapeOf`, meanwhile, - // is [OUT, IN], so the same two tensors are demanded below as - // {ple_dim, hidden} and {hidden, ple_dim}. - // - // The shapes are what stop the two being swapped: they have the - // SAME element count on this model (1536 x 256 against - // 256 x 1536), so no length check anywhere would notice, and - // "passing one the other way round packs silently and produces - // noise". - // - // THE MAPPING IS THE TRANSPOSE OF WHAT THE PACKER WANTS, and the - // pack lambda below is where that is undone. `inp_gate` maps as - // ple_dim contiguous rows of hidden and corelib wants hidden rows - // of ple_dim; `proj` maps as hidden rows of ple_dim and corelib - // wants ple_dim rows of hidden. Both are the transpose of the - // other side, and the element counts are IDENTICAL, so nothing - // rejects the wrong one -- "passing one the other way round packs - // silently and produces noise". The reference driver does exactly - // the same thing at exactly this point - // (`gemma4_driver.py:1517`: `np.ascontiguousarray(floats(...))`, - // where `floats()` is `.reshape(reversed(dims)).T`). - // - // These two views are the MAPPING; they are never handed to the - // packer's filling leg unmodified. - inp_gate[index] = f32_2d(Blk(layer, ".inp_gate.weight"), cfg.ple_dim, cfg.hidden); - proj[index] = f32_2d(Blk(layer, ".proj.weight"), cfg.hidden, cfg.ple_dim); - ple_post_norm[index] = f32_1d(Blk(layer, ".post_norm.weight"), cfg.hidden); - // THE NEXT LAYER'S ATTENTION NORM, AND THE LAST LAYER'S IS THE - // MODEL'S FINAL NORM. `ple` fuses [gate -> proj -> norm -> - // residual -> scale -> norm] and that last norm crosses the block - // boundary: it is the NEXT block's input norm. There is no - // `blk.{layers}.attn_norm`, so the last layer takes - // `output_norm.weight`. The driver's rule verbatim - // (`_next_norm_name`), and the header's statement of the stakes: - // getting it off by one "is not a crash -- every layer is then - // normalized by its neighbour's gamma, which still generates - // text." - ple_next_norm[index] = f32_1d( - layer == cfg.layers - 1 ? std::string("output_norm.weight") - : Blk(layer + 1, ".attn_norm.weight"), - cfg.hidden); - // `layer_output_scale` is a ONE-ELEMENT tensor -- a norm's shape - // without being one -- and is the (x + o) multiplier baked into - // the blob. - layer_scale[index] = f32_1d(Blk(layer, ".layer_output_scale.weight"), 1).values[0]; - } - - // ---- the stream, then the plan, then the weights - // - // In that order and not another: corelib 0.5.0 takes the stream on - // every padding helper (the PDI pair a stream was opened with selects - // the ELF set, and a shape exists under one pair and not another), so - // the plan cannot exist before a stream does. - auto lease = runtime->AcquireExecution(); - const int token_pdi = paged ? kPagedTokenPdi : kTokenPdi; - stream = corelib::OpenModelStream(api, kPrefillPdi, token_pdi, "gemma4"); - plan = Gemma4ShapePlan::Build(api, stream.get(), cfg, - paged ? static_cast(maximum) - : kMaxSequenceLength); - - // ---- the slot table - std::vector slots; - slots.reserve(layer_count * 10 + 2); - - const auto add_matmul = [&](std::string label, std::int64_t k, std::int64_t n, - std::uint32_t group, const GgufTensor* blocks, - UniqueMatMulWeights* destination) { - WeightSlot slot; - slot.kind = SlotKind::Matmul; - slot.label = std::move(label); - slot.matmul = {k, n, group, /*has_bias=*/false}; - const auto desc = slot.matmul; - const auto name = slot.label; - slot.pack = [this, desc, name, blocks, destination] { - const ryzenai_corelib_weights_source sources[]{corelib::GgufQ8( - ryzenai_corelib_weights_role_qweight, blocks->bytes.data(), - blocks->bytes.size(), desc.n, desc.k)}; - const auto options = corelib::PackOptions(kRequantizeThreads); - void* created = nullptr; - api->Check(api->functions().matmul_weights_create( - &desc, sources, std::size(sources), &options, &created), - "ryzenai_corelib_matmul_bf16_weights_create " + name); - *destination = UniqueMatMulWeights(api, created); - }; - slot.assign = [this, destination](void* handle) { - *destination = UniqueMatMulWeights(api, handle); - }; - slot.handle = [destination] { return destination->get(); }; - slots.push_back(std::move(slot)); - }; - - const auto add_rmsnorm = [&](std::string label, std::int64_t k, - const std::vector* gamma, - UniqueRmsNormWeights* destination) { - WeightSlot slot; - slot.kind = SlotKind::RmsNorm; - slot.label = std::move(label); - // `epsilon` travels in the DESCRIPTOR here, not per dispatch: it - // is baked into the packed blob. There is no group size and no - // layout to choose -- the scale is BF16 and unquantized -- and - // `k` must match a shipped artifact exactly, because this op - // tiles the ROW count and cannot pad k. - slot.rmsnorm = {k, cfg.eps}; - const auto desc = slot.rmsnorm; - const auto name = slot.label; - slot.pack = [this, desc, name, gamma, destination] { - const ryzenai_corelib_weights_source sources[]{corelib::Bf16( - ryzenai_corelib_weights_role_scale, gamma->data(), gamma->size())}; - void* created = nullptr; - api->Check(api->functions().rmsnorm_weights_create( - &desc, sources, std::size(sources), nullptr, &created), - "ryzenai_corelib_rmsnorm_bf16_weights_create " + name); - *destination = UniqueRmsNormWeights(api, created); - }; - slot.assign = [this, destination](void* handle) { - *destination = UniqueRmsNormWeights(api, handle); - }; - slot.handle = [destination] { return destination->get(); }; - slots.push_back(std::move(slot)); - }; - - // LAYER 0'S STANDALONE NORM COMES FIRST, matching the reference - // driver's own order (`_load_layers` packs it before the layers). - add_rmsnorm("blk.0.attn_norm.weight", cfg.hidden, &input_norm_bf16, &input_norm); - - for (std::size_t index = 0; index < layer_count; ++index) { - const auto layer = static_cast(index); - const bool owns = layer < cfg.kv_layers; - const auto head_dim = cfg.LayerHeadDim(layer); - const auto q_dim = cfg.q_heads * head_dim; - const auto kv_dim = cfg.kv_heads * head_dim; - const auto intermediate = cfg.layer_intermediate[index]; - auto& weights = layers[index]; - - // EVERY PROJECTION IS `cfg.group`. Not a judgement call: the - // driver declares the field as "group: int # every projection", - // `_matmul_weights` defaults `group_size=self.cfg.group`, and the - // ONLY call in the whole driver that overrides it is lm_head. - // `cfg.head_group` is named for lm_head, NOT for the attention - // heads -- using it here would requantize Q/K/V/O at the wrong - // group, which is not a padding question but a wrong weight. It - // happens to be the same number on both shipped rows, which is - // exactly why it has to be written down. - add_matmul(Blk(layer, ".attn_q.weight"), cfg.hidden, q_dim, cfg.group, - &q_weight[index], &weights.q); - if (owns) { - add_matmul(Blk(layer, ".attn_k.weight"), cfg.hidden, kv_dim, cfg.group, - &k_weight[index], &weights.k); - add_matmul(Blk(layer, ".attn_v.weight"), cfg.hidden, kv_dim, cfg.group, - &v_weight[index], &weights.v); - } - add_rmsnorm(Blk(layer, ".attn_q_norm.weight"), head_dim, &q_norm[index], - &weights.q_norm); - if (owns) { - add_rmsnorm(Blk(layer, ".attn_k_norm.weight"), head_dim, &k_norm[index], - &weights.k_norm); - // V's norm, which has NO learned weight in the file. - add_rmsnorm(Blk(layer, ".attn_v_norm(ones)"), head_dim, - head_dim == cfg.head_dim ? &ones_sliding : &ones_full, - &weights.v_norm); - } - add_matmul(Blk(layer, ".attn_output.weight"), q_dim, cfg.hidden, cfg.group, - &o_weight[index], &weights.o); - add_rmsnorm(Blk(layer, ".post_attention_norm.weight"), cfg.hidden, - &post_attn_norm[index], &weights.post_attn_norm); - - { - WeightSlot slot; - slot.kind = SlotKind::SsMlp; - slot.label = Blk(layer, ".ssmlp"); - slot.layer = layer; - slot.ssmlp = {cfg.hidden, intermediate, cfg.group, - /*activation=*/1, /*post_feedforward_layernorm=*/1}; - const auto desc = slot.ssmlp; - const auto name = slot.label; - const auto* norm0 = &ffn_norm[index]; - const auto* norm1 = &post_ffw_norm[index]; - const auto* gate = &gate_weight[index]; - const auto* up = &up_weight[index]; - const auto* down = &down_weight[index]; - const auto* eps = ε - auto* destination = &weights.mlp; - slot.pack = [this, desc, name, norm0, norm1, gate, up, down, eps, - destination] { - const ryzenai_corelib_weights_source sources[]{ - corelib::Bf16(ryzenai_corelib_weights_role_epsilon, eps->data(), eps->size()), - corelib::Bf16(ryzenai_corelib_weights_role_norm0, norm0->data(), norm0->size()), - corelib::Bf16(ryzenai_corelib_weights_role_norm1, norm1->data(), norm1->size()), - corelib::GgufQ8(ryzenai_corelib_weights_role_gate_qweight, - gate->bytes.data(), gate->bytes.size(), desc.n, desc.k), - corelib::GgufQ8(ryzenai_corelib_weights_role_up_qweight, - up->bytes.data(), up->bytes.size(), desc.n, desc.k), - corelib::GgufQ8(ryzenai_corelib_weights_role_down_qweight, - down->bytes.data(), down->bytes.size(), desc.k, desc.n)}; - const auto options = corelib::PackOptions(kRequantizeThreads); - void* created = nullptr; - api->Check(api->functions().ssmlp_weights_create( - &desc, sources, std::size(sources), &options, &created), - "ryzenai_corelib_ssmlp_bf16_weights_create " + name); - *destination = UniqueSsMlpWeights(api, created); - }; - slot.assign = [this, destination](void* handle) { - *destination = UniqueSsMlpWeights(api, handle); - }; - slot.handle = [destination] { return destination->get(); }; - slots.push_back(std::move(slot)); - } - - { - WeightSlot slot; - slot.kind = SlotKind::Ple; - slot.label = Blk(layer, ".ple"); - slot.layer = layer; - // `group_size` here is `ple_group`, NOT `group`: 32 is the - // only value mladfple ships and corelib refuses anything else - // by name. `epsilon` is ONE value added under the rsqrt of - // BOTH RMSNorms, not one per norm. - slot.ple = {cfg.hidden, cfg.ple_dim, cfg.ple_group, cfg.eps, - layer_scale[index]}; - const auto desc = slot.ple; - const auto name = slot.label; - const auto* gate = &inp_gate[index]; - const auto* projection = &proj[index]; - const auto* post_norm = &ple_post_norm[index]; - const auto* next_norm = &ple_next_norm[index]; - auto* destination = &weights.ple; - slot.pack = [this, desc, name, gate, projection, post_norm, next_norm, - destination] { - // corelib takes gate as FP32 [k, n] and proj as [n, k]. - // The GGUF mapping is the other way round, so transpose - // first. The packer allocates its own blob. - const auto gate_rows = Transposed(gate->values, desc.n, desc.k); - const auto proj_rows = Transposed(projection->values, desc.k, desc.n); - const ryzenai_corelib_weights_source sources[]{ - corelib::Fp32(ryzenai_corelib_weights_role_gate, gate_rows.data(), - gate_rows.size(), {desc.k, desc.n}), - corelib::Fp32(ryzenai_corelib_weights_role_proj, proj_rows.data(), - proj_rows.size(), {desc.n, desc.k}), - corelib::Fp32(ryzenai_corelib_weights_role_post_norm, - post_norm->values.data(), post_norm->values.size(), - {desc.k}), - corelib::Fp32(ryzenai_corelib_weights_role_next_norm, - next_norm->values.data(), next_norm->values.size(), - {desc.k})}; - void* created = nullptr; - api->Check(api->functions().ple_weights_create( - &desc, sources, std::size(sources), nullptr, &created), - "ryzenai_corelib_ple_bf16_weights_create " + name); - *destination = UniquePleWeights(api, created); - }; - slot.assign = [this, destination](void* handle) { - *destination = UniquePleWeights(api, handle); - }; - slot.handle = [destination] { return destination->get(); }; - slots.push_back(std::move(slot)); - } - } - - // LM_HEAD, LAST, AND THE ONE WEIGHT AT `head_group`. TIED to - // token_embd -- this conversion has no `output.weight` on either row - // -- and "the only weight packed at cfg.head_group" - // (gemma4_driver.py:1691-1694, :2449-2453). - add_matmul("token_embd.weight", cfg.hidden, cfg.vocab, cfg.head_group, - &embedding, &lm_head); - slot_count = slots.size(); - - // A miss, a stale key, an unreadable index or a failed bind all fall - // back to packing. - const auto cache = WeightCache::ForGguf(package->Path()); - const auto version = api->runtime_version(); - const auto cache_key = MakeWeightCacheKey( - package->Path(), version.major, version.minor, version.patch, - CacheLayout(cfg), slots.size()); - if (cache) { - if (const auto spans = cache->ReadIndex(cache_key)) - from_cache = BindFromCache(cache->DataPath(), *spans, slots); - } - - if (!from_cache) { - std::cout << "[FLM] Packing " << slots.size() << " Gemma 4 " << cfg.row_name - << " weights from Q8_0" << (cache ? "; cached for the next load" : "") - << std::endl; - PackAll(slots); - } - - if (cache && !from_cache) { - // Whatever is there did not match, or it would have been used. - (void)cache->Remove(); - WriteCache(*cache, cache_key, slots); - } - - AllocateBuffers(); - AllocateRotaryTables(rope_factors); - } - - /// \brief every activation buffer, the head's two, the per-layer embedding - /// planes and the KV caches -- once, at load - /// \note See the member declarations for why each is shaped as it is. The - /// order here is the order the shapes are asserted in. - void AllocateBuffers() { - const auto& cfg = shape; - const auto rows = kMaxSequenceLength; - const auto widest = std::max(cfg.head_dim, cfg.global_head_dim); - const auto q_dim = cfg.q_heads * widest; - const auto kv_dim = cfg.kv_heads * widest; - const auto bf16 = [&](std::vector dims, const char* label) { - return Tensor(ryzenai_corelib_data_type_bf16, std::move(dims), label); - }; - - hidden = bf16({rows, cfg.hidden}, "hidden"); - residual = bf16({rows, cfg.hidden}, "residual"); - skip_sum = bf16({rows, cfg.hidden}, "skip_sum"); - ple_out = bf16({rows, cfg.hidden}, "ple_out"); - q_buf = bf16({rows, q_dim}, "q"); - k_buf = bf16({rows, kv_dim}, "k"); - v_buf = bf16({rows, kv_dim}, "v"); - q_normed = bf16({rows, q_dim}, "q normed"); - k_normed = bf16({rows, kv_dim}, "k normed"); - attn_out = bf16({rows, q_dim}, "attention"); - last_hidden = bf16({1, cfg.hidden}, "lm head input"); - logits = bf16({1, cfg.vocab}, "logits"); - - ple_in.reserve(static_cast(cfg.layers)); - for (std::int64_t layer = 0; layer < cfg.layers; ++layer) - ple_in.push_back(bf16({rows, cfg.ple_dim}, "per-layer embedding plane")); - - k_cache.reserve(static_cast(cfg.kv_layers)); - v_cache.reserve(static_cast(cfg.kv_layers)); - for (std::int64_t cache = 0; cache < cfg.kv_layers; ++cache) { - // THE WIDTH IS THE OWNING LAYER'S OWN HEAD. A layer below - // `kv_layers` owns cache `index`, so this is layer `cache`'s head. - const auto head = cfg.LayerHeadDim(cache); - if (paged) { - const auto pages = static_cast(cache_rows) / kPageSize; - k_cache.push_back(bf16({pages, cfg.kv_heads, kPageSize, head}, "K cache")); - v_cache.push_back(bf16({pages, cfg.kv_heads, kPageSize, head}, "V cache")); - } else { - k_cache.push_back(bf16({cfg.kv_heads, rows, head}, "K cache")); - v_cache.push_back(bf16({cfg.kv_heads, rows, head}, "V cache")); - } - } - } - - /// \brief create one device tensor, naming it in any failure - UniqueTensor Tensor(ryzenai_corelib_data_type type, - std::vector dims, const char* label) { - void* created = nullptr; - api->Check(api->functions().create_device_tensor(type, dims.data(), dims.size(), - &created), - std::string("ryzenai_corelib_create_device_tensor ") + label); - return UniqueTensor(api, created); - } - - /// \brief both rotary pairs, built once and never written again - /// \param factors `rope_freqs.weight`, the FULL-ATTENTION regime's - /// per-pair frequency divisors - /// - /// SLIDING LAYERS GET NO FACTORS AND FULL-ATTENTION LAYERS DO, which is - /// llama.cpp's own split: `freq_factors` is passed on full-attention - /// layers only. Applying them to both regimes, or to neither, produces a - /// perfectly ordinary rotary table that is wrong from position 1 on and - /// still generates fluent text. - /// - /// Sized at `kMaxSequenceLength` rather than at the prompt: corelib - /// requires both tables to "reach at least `position + rows` rows", and - /// they are indexed by ABSOLUTE position, so the table a decode step at - /// position 4000 needs is the same table prefill wrote. - void AllocateRotaryTables(const GgufFloatTensor& factors) { - const auto& cfg = shape; - const auto view = [&](UniqueHostView& destination, std::int64_t head_dim, - const std::vector& values, const char* label) { - const std::array dims{static_cast(cache_rows), - head_dim / 2}; - void* created = nullptr; - api->Check(api->functions().create_host_view( - ryzenai_corelib_data_type_fp32, dims.data(), dims.size(), - values.data(), &created), - std::string("ryzenai_corelib_create_host_view ") + label); - destination = UniqueHostView(api, created); - }; - // Base `rope_theta_swa` (1e4) over the sliding head, NO factors. - const auto rope_rows = static_cast(cache_rows); - auto sliding = MakeRopeTables(cfg.head_dim, cfg.rope_theta_swa, rope_rows); - rope_cos_sliding = std::move(sliding.cos); - rope_sin_sliding = std::move(sliding.sin); - view(cos_sliding, cfg.head_dim, rope_cos_sliding, "sliding cos"); - view(sin_sliding, cfg.head_dim, rope_sin_sliding, "sliding sin"); - // Base `rope_theta` (1e6) over the full head, WITH `rope_freqs.weight` - // as a per-pair frequency DIVISOR. - auto full = MakeRopeTables(cfg.global_head_dim, cfg.rope_theta, rope_rows, - factors.values); - rope_cos_full = std::move(full.cos); - rope_sin_full = std::move(full.sin); - view(cos_full, cfg.global_head_dim, rope_cos_full, "full cos"); - view(sin_full, cfg.global_head_dim, rope_sin_full, "full sin"); - } - - /// \brief run every slot's `pack()` across a small pool - /// \note TWO FORMS OF PARALLELISM THAT MULTIPLY. Each create is - /// submitted at `kRequantizeThreads` (0 -- corelib picks, 8 packer - /// threads at 0.25.0) while `kWeightCreateConcurrency` of them run - /// at once, so up to the product runs together. The reference - /// driver packs strictly serially and says why (an open, - /// unattributed corelib defect correlating concurrent creates with - /// all-zero device output); this follows Phi-4's shape instead, - /// which is what this tree already ships and runs. - void PackAll(std::vector& slots) { - std::atomic next{0}; - std::mutex failure_mutex; - std::exception_ptr first_failure; - const auto run = [&] { - for (auto index = next++; index < slots.size(); index = next++) { - try { - slots[index].pack(); - } catch (...) { - std::lock_guard lock(failure_mutex); - if (!first_failure) first_failure = std::current_exception(); - // Stop the others: the first diagnostic is the useful - // one, and a package that fails one create fails the load. - next = slots.size(); - return; - } - } - }; - const auto workers = std::min(kWeightCreateConcurrency, slots.size()); - std::vector pool; - pool.reserve(workers > 0 ? workers - 1 : 0); - for (std::size_t worker = 1; worker < workers; ++worker) pool.emplace_back(run); - run(); - for (auto& worker : pool) worker.join(); - if (first_failure) std::rethrow_exception(first_failure); - } - - /// \brief bind every weight from the cache file instead of packing - /// \return true when ALL of them bound; false leaves nothing bound - /// \note A PARTIAL BIND RELEASES EVERYTHING. A half-cached model is not a - /// state worth having and the caller simply packs. corelib rejects - /// a slice whose length is not exactly what the descriptor packs - /// to, so a stale file is caught here rather than becoming - /// confident nonsense. - bool BindFromCache(const std::filesystem::path& data_path, - const std::vector& spans, - std::vector& slots) { - if (spans.size() != slots.size()) return false; - const auto path = data_path.string(); - for (std::size_t index = 0; index < slots.size(); ++index) { - auto& slot = slots[index]; - void* handle = nullptr; - auto status = ryzenai_corelib_status_failure; - const auto packed = - corelib::PackedFile(path.c_str(), spans[index].offset, spans[index].size); - switch (slot.kind) { - case SlotKind::Matmul: - status = api->functions().matmul_weights_create( - &slot.matmul, &packed, 1, nullptr, &handle); - break; - case SlotKind::SsMlp: - status = api->functions().ssmlp_weights_create( - &slot.ssmlp, &packed, 1, nullptr, &handle); - break; - case SlotKind::RmsNorm: - status = api->functions().rmsnorm_weights_create( - &slot.rmsnorm, &packed, 1, nullptr, &handle); - break; - case SlotKind::Ple: - status = api->functions().ple_weights_create( - &slot.ple, &packed, 1, nullptr, &handle); - break; - } - if (status != ryzenai_corelib_status_success || handle == nullptr) { - ReleaseAllWeights(); - return false; - } - slot.assign(handle); - } - return true; - } - - void ReleaseAllWeights() { - for (auto& layer : layers) layer = LayerWeights{}; - input_norm = {}; - lm_head = {}; - } - - /// \brief copy the packed bytes out and write them beside the model - /// \note Best effort: a cache that cannot be written is not a load - /// failure, it only means the next launch packs again. The index is - /// written last, so a data file without a matching index is never - /// used. - void WriteCache(const WeightCache& cache, const WeightCacheKey& key, - const std::vector& slots) { - try { - std::error_code error; - const auto data_path = cache.DataPath(); - std::filesystem::create_directories(data_path.parent_path(), error); - const auto temporary = data_path.string() + ".tmp"; - std::vector spans; - spans.reserve(slots.size()); - { - std::ofstream output(temporary, std::ios::binary | std::ios::trunc); - if (!output) return; - std::vector buffer; - std::uint64_t offset = 0; - for (const auto& slot : slots) { - std::size_t size = 0; - if (api->functions().weights_copy_data(slot.handle(), nullptr, 0, &size) != - ryzenai_corelib_status_success || - size == 0) - return; - buffer.resize(size); - if (api->functions().weights_copy_data( - slot.handle(), buffer.data(), buffer.size(), &size) != - ryzenai_corelib_status_success) - return; - if (size == 0) return; - output.write(buffer.data(), static_cast(size)); - if (!output) return; - spans.push_back({offset, static_cast(size)}); - offset += size; - } - } - std::filesystem::rename(temporary, data_path, error); - if (error) { std::filesystem::remove(temporary, error); return; } - if (!cache.WriteIndex(key, spans)) std::filesystem::remove(data_path, error); - } catch (...) { - // Caching is an optimisation; never let it fail a load. - } - } - - void usable() const { - if (poisoned) throw std::runtime_error("Gemma 4 corelib engine is poisoned"); - } - - // ---- the forward pass ------------------------------------------------- - - /// \brief create a window onto `parent`, naming it in any failure - UniqueTensorWindow Window(void* parent, std::vector dims, - std::size_t offset, const char* label, - std::vector strides = {}) { - void* created = nullptr; - api->Check(api->functions().create_tensor_window(parent, dims.data(), dims.size(), offset, - strides.empty() ? nullptr : strides.data(), - strides.size(), &created), - std::string("ryzenai_corelib_create_tensor_window ") + label); - return UniqueTensorWindow(api, created); - } - - /// \brief the row count the kernels will actually run for `rows` live ones - /// - /// THE GRID STARTS AT 128 BECAUSE PRODUCTION'S DOES. A four-token prompt - /// runs a 128-row kernel here where every other driver in this tree would - /// run a 64-row one. That is not a mitigation for a broken artifact: 128 - /// is the floor hybrid-llm's full-fusion path supports at all, and below - /// it upstream's own comment says a native buffer wedges the NPU for LLM - /// matmul shapes. The reference driver's `PREFILL_GRID` is this table and - /// `_pad_rows` is this function. The 124 rows of zeros are what staying on - /// a supported shape costs. - /// - /// DECODE IS NOT PADDED: an M == 1 kernel ships for every operator here. - std::int64_t PaddedRowsFor(std::int64_t rows) const { - if (rows <= 1) return 1; - // A paged chunk past page 0 is only correct at a full page. - if (paged && position > 0 && rows <= kPageSize) return kPageSize; - for (const std::int64_t bucket : {128, 256, 512, 1024, 2048, 3072, 4096}) { - if (rows <= bucket) return bucket; - } - throw std::out_of_range("Gemma 4 cannot run " + std::to_string(rows) + - " rows: the largest shipped prefill kernel is 4096"); - } - - /// \brief byte-element offset of one KV row inside its cache - std::size_t CacheOffset(std::int64_t kv_head, std::int64_t index, - std::int64_t head) const { - if (!paged) { - return static_cast((kv_head * kMaxSequenceLength + index) * head); - } - const auto page = index / kPageSize; - const auto row = index % kPageSize; - return static_cast( - ((page * shape.kv_heads + kv_head) * kPageSize + row) * head); - } - - /// \brief every operator agrees that `padded` is the extent it runs at - /// - /// ONE ROW EXTENT FOR THE WHOLE PASS, AND IT HAS TO BE CHECKED RATHER THAN - /// ASSUMED. Every window this pass builds declares `padded` rows, and each - /// operand "is checked against its OWN extent" by corelib -- so if any - /// operator's own pad helper would round `padded` further, that operator - /// reads and writes rows nobody allocated for it. The plan already asked - /// corelib, for every bucket and both geometries, at load; this compares - /// the answers. - /// - /// CHECKED ON THE HOST, BEFORE THE FIRST DISPATCH, for the reason every - /// other check in `Run` is: a rejection part-way through an unbarriered - /// chain aborts the process rather than raising. - /// - /// WHAT THIS DOES **NOT** COVER, AND THE REFERENCE DRIVER DOES. - /// `gemma4_driver.py` carries a separate `_norm_rows` guard asserting that - /// each RMSNorm runs at a row count an EXACT kernel ships for. It exists - /// because an unshipped M does not fail -- it OVERSHOOTS, and for V's norm - /// the destination is the KV CACHE, so the overshoot lands in another - /// layer's cached values and comes back as attention over data no layer - /// wrote. Nothing here reproduces that guard. - /// - /// It is latent today ONLY because - /// `ryzenai_corelib_rmsnorm_bf16_pad_rows` is an identity function that - /// does not report the covering extent (corelib.h says so outright) and is - /// therefore not in `FLM_CORELIB_FUNCTIONS` at all -- so the rmsnorm rows - /// this engine uses are `ple`'s bucket, which is a shipped one. That is a - /// reason that can stop being true the day corelib implements the helper, - /// and the failure it would unmask is silent. Carried as a pre-hardware - /// debt item; see the SDD ledger's pre-hardware debt list. - void RequireOneRowExtent(std::int64_t padded) const { - const auto for_geometry = [&](bool swa) { - const auto& kinds = shape.layer_is_swa; - if (std::find(kinds.begin(), kinds.end(), swa) == kinds.end()) return; - const auto& extents = plan.ForRows(static_cast(padded), swa); - const std::array, 6> named{{ - {"query", extents.q_rows}, - {"key/value", extents.kv_rows}, - {"output", extents.output_rows}, - {"ssmlp", extents.ssmlp_rows}, - {"ple", extents.ple_rows}, - {"attention", extents.flat_mha_rows}, - }}; - for (const auto& [what, value] : named) { - if (value == padded) continue; - throw std::runtime_error( - std::string("Gemma 4 would run ") + std::to_string(padded) + - " rows, but corelib pads " + what + " to " + std::to_string(value) + - " on the " + (swa ? "sliding" : "full-attention") + - " geometry -- every buffer in this pass is windowed at " + - std::to_string(padded) + " rows and that op would run past it"); - } - }; - for_geometry(true); - for_geometry(false); - } - - /// \brief the windows one geometry's layers dispatch through - /// \note THE PER-HEAD VIEWS ARE WINDOWS OF THE SAME MEMORY, not copies. - /// Reading `[rows, heads * head]` as `[rows * heads, head]` IS the - /// per-head layout in row-major, because a token's heads are - /// adjacent -- which is why Q and K need no staging buffer and V - /// does: the CACHE is the other way round. - struct Geometry { - std::int64_t head{}; - UniqueTensorWindow q, k, v_dest; - UniqueTensorWindow q_heads_in, k_heads_in, q_heads_out, k_heads_out; - std::vector v_src; - UniqueTensorWindow q_attn, k_attn, attn; - }; - - Geometry BuildGeometry(bool swa, std::int64_t padded) { - const auto& cfg = shape; - Geometry geometry; - const auto head = swa ? cfg.head_dim : cfg.global_head_dim; - geometry.head = head; - const auto q_dim = cfg.q_heads * head; - const auto kv_dim = cfg.kv_heads * head; - geometry.q = Window(q_buf.get(), {padded, q_dim}, 0, "q"); - geometry.k = Window(k_buf.get(), {padded, kv_dim}, 0, "k"); - // WHERE V'S PROJECTION LANDS, AND IT IS HEAD-MAJOR BECAUSE THE CACHE - // IS. A cache is [kv_heads, max_seq, head]; a projection's ordinary - // output is [rows, kv_heads * head], token-major. Those are the same - // bytes only when there is ONE KV head. With more, one head's rows sit - // `kv_heads * head` apart in the token-major output, while V's norm - // below reads each head as a dense [rows, head] plane at an offset -- - // so the staging buffer has to be written head-major in the first - // place. Only the scattering kernel can do that, and a 3-D output, a - // window strided [max_seq * head, head, 1], is what asks for it. - // - // AT EXACTLY ONE KV HEAD A PLAIN 2-D OUTPUT IS USED INSTEAD, and that - // is not merely an optimization: with one head there is nothing to - // scatter and the dense kernel writes byte-for-byte what the - // scattering one would, while the gm artifact at one head and head - // size 512 TAKES THE DEVICE DOWN (`gemma4_driver.py`'s - // GM_AT_ONE_KV_HEAD). The rule is `cfg.kv_heads`, not a model. - geometry.v_dest = - cfg.kv_heads == 1 - ? Window(v_buf.get(), {padded, head}, 0, "v") - : Window(v_buf.get(), {cfg.kv_heads, kMaxSequenceLength, head}, 0, "v", - {kMaxSequenceLength * head, head, 1}); - // QK-Norm normalizes ONE HEAD, so its row count is tokens x heads. - geometry.q_heads_in = - Window(q_buf.get(), {padded * cfg.q_heads, head}, 0, "q heads"); - geometry.k_heads_in = - Window(k_buf.get(), {padded * cfg.kv_heads, head}, 0, "k heads"); - geometry.q_heads_out = - Window(q_normed.get(), {padded * cfg.q_heads, head}, 0, "q heads normed"); - geometry.k_heads_out = - Window(k_normed.get(), {padded * cfg.kv_heads, head}, 0, "k heads normed"); - // V's norm reads one head's plane at a time. The staging buffer - // carries `max_seq` rows per head, exactly as the cache does, so head - // `h` starts at the same offset in both. - geometry.v_src.reserve(static_cast(cfg.kv_heads)); - for (std::int64_t kv_head = 0; kv_head < cfg.kv_heads; ++kv_head) { - geometry.v_src.push_back(Window( - v_buf.get(), {padded, head}, - static_cast(kv_head * kMaxSequenceLength * head), - "v head")); - } - geometry.q_attn = Window(q_normed.get(), {padded, q_dim}, 0, "q attended"); - geometry.k_attn = Window(k_normed.get(), {padded, kv_dim}, 0, "k attended"); - geometry.attn = Window(attn_out.get(), {padded, q_dim}, 0, "attention"); - return geometry; - } - - /// \brief one token's embedding row, scaled by `sqrt(hidden)` - /// - /// THE SCALE IS PART OF THE EMBEDDING, as it is in llama.cpp: - /// `ggml_scale(inpL, sqrtf(n_embd))` runs immediately after the gather and - /// before anything reads it. So the per-layer inputs must NOT scale again; - /// a second `sqrt(1536)` there multiplies the projection by 39 and still - /// produces finite, fluent-looking output. - /// - /// `token_embd.weight`'s file dims are `[hidden, vocab]` with the input - /// width fastest-varying, so the memory is `vocab` rows of `hidden` and - /// one token's row is contiguous -- the same layout question - /// `GatherPerLayerEmbedding` answers for the per-layer table. - std::vector ScaledEmbeddingRows(std::span ids) const { - const auto& cfg = shape; - constexpr std::int64_t kBlockElements = 32; - constexpr std::size_t kBlockBytes = 34; - if (cfg.hidden % kBlockElements != 0) { - throw std::runtime_error("hidden is " + std::to_string(cfg.hidden) + - ", which is not a multiple of the Q8_0 block size"); - } - const auto row_bytes = - static_cast(cfg.hidden / kBlockElements) * kBlockBytes; - if (embedding_table.bytes.size() < - static_cast(cfg.vocab) * row_bytes) { - throw std::runtime_error("token_embd.weight is shorter than " + - std::to_string(cfg.vocab) + " rows of " + - std::to_string(cfg.hidden)); - } - const auto scale = static_cast(std::sqrt(static_cast(cfg.hidden))); - std::vector rows(ids.size() * static_cast(cfg.hidden)); - for (std::size_t row = 0; row < ids.size(); ++row) { - const auto span = std::span(rows).subspan( - row * static_cast(cfg.hidden), - static_cast(cfg.hidden)); - DecodeQ8Row(embedding_table.bytes.subspan( - static_cast(ids[row]) * row_bytes, row_bytes), - span); - for (float& value : span) value *= scale; - } - return rows; - } - - /// \brief one pass over every layer, then the head. Returns SOFTCAPPED - /// logits for the LAST position. - /// - /// NO LAYER BOUNDARY RETURNS TO THE HOST, which is the whole point of - /// `ple_bf16`: its SECOND output is the next layer's input norm, so the - /// chain closes on the device and the 306 dispatches of an E2B prefill - /// carry nothing back here. - /// - /// ONE WAIT IN THE LOOP, and it is after the loop. There is NO BARRIER AT - /// A ROTARY REGIME CHANGE: corelib holds one derived table per rotary key - /// per synchronize interval and never rewrites one the interval has - /// already bound, so "no barrier is needed at a change of base". The - /// reference driver removed exactly that barrier and records what it cost: - /// 13 extra round trips on an E2B. - /// - /// EVERY ARGUMENT IS CHECKED BEFORE THE FIRST DISPATCH, AND THAT IS A RULE - /// OF THE WHOLE DISPATCH MODEL RATHER THAN A HOUSE STYLE. Any op rejection - /// part-way through an unbarriered chain is PROCESS-FATAL, whatever the op - /// and whatever the reason: the exception unwinds out of here without - /// reaching `stream_synchronize`, so the runs already submitted are - /// destructed while their commands are still in flight and XRT aborts -- - /// "xrt::run destructed while command is still in progress" -- with no - /// catchable failure anywhere. There is no try/finally that fixes it - /// either: synchronizing during the unwind waits on work whose operands - /// are already being torn down. So anything this pass can refuse, it - /// refuses HERE, before the stream has been handed anything. - buffer Run(std::span ids) { - usable(); - const auto& cfg = shape; - if (ids.empty()) - throw std::invalid_argument("Gemma 4 request contains no token IDs"); - for (const int id : ids) { - if (id < 0 || id >= cfg.vocab) { - throw std::out_of_range( - "Gemma 4 token id " + std::to_string(id) + - " is outside the vocabulary of " + std::to_string(cfg.vocab)); - } - } - const auto rows = static_cast(ids.size()); - if (rows > static_cast(max_length) || - position + rows > static_cast(max_length)) { - throw std::out_of_range("Gemma 4 request exceeds configured context capacity"); - } - const auto padded = PaddedRowsFor(rows); - // The kernels run `padded` rows starting at `position`, and V's norm - // writes that many cache rows -- so the bound is on `position + - // padded`, not on `position`. Reported here, where the numbers are - // still the caller's, rather than as a byte offset from inside a - // window. - // - // UNREACHABLE TODAY, AND SAID SO RATHER THAN LEFT TO LOOK LIVE. The - // check above bounds `position + rows` by `max_length`, which the - // constructor bounds by `kMaxSequenceLength`; `padded` exceeds `rows` - // only when `rows > 1`, which the chunked-prefill rule below pins to - // `position == 0`, and `PaddedRowsFor` refuses anything over 4096. So - // no argument reaches this branch while `max_length <= - // kMaxSequenceLength` holds, and no test in this suite can reach it. - // It is kept because it guards the CACHE and the one above guards the - // CONFIGURATION: they are independent facts that happen to coincide, - // the reference driver has this one and not the other, and a future - // caller that decouples them would otherwise lose the only check that - // is about the thing the kernel actually addresses. - const auto cache_limit = static_cast(cache_rows); - if (position < 0 || position + padded > cache_limit) { - throw std::out_of_range( - std::to_string(rows) + " rows at position " + std::to_string(position) + - " run a " + std::to_string(padded) + "-row kernel, which reaches row " + - std::to_string(position + padded) + " of a " + - std::to_string(cache_limit) + "-row cache"); - } - // Chunked prefill is paged only, and only on a page boundary. - // flat_mha rejects a multi-row call at a non-zero position itself, but - // as the sixth dispatch of a layer, which aborts the process. - const bool page_chunk = paged && position > 0 && position % kPageSize == 0; - if (rows > 1 && position != 0 && !page_chunk) { - throw std::invalid_argument( - std::to_string(rows) + " rows at position " + std::to_string(position) + - ": a multi-row pass must start at position 0" + - (paged ? " or on a 4096-row page boundary" : "")); - } - RequireOneRowExtent(padded); - - // ---- the host's share, all of it before anything is dispatched - const auto embedded = ScaledEmbeddingRows(ids); - const auto planes = BuildPerLayerInputs( - ple_table.bytes, ple_projection.values, ple_projection_norm.values, - embedded, ids, cfg.layers, cfg.ple_dim, cfg.hidden, cfg.vocab, cfg.eps); - - auto lease = runtime->AcquireExecution(); - // THE RESIDUAL STREAM ENTERING LAYER 0 IS THE SCALED EMBEDDING, and it - // is also what layer 0's norm reads -- so ONE staged write, and the - // norm below writes `hidden` entirely. FP32 into a BF16 tensor: - // corelib converts on the way in. - std::vector staged(static_cast(padded * cfg.hidden), 0.0f); - std::copy(embedded.begin(), embedded.end(), staged.begin()); - api->Check(api->functions().tensor_write(residual.get(), - ryzenai_corelib_data_type_fp32, - staged.data(), staged.size(), 0), - "ryzenai_corelib_tensor_write residual"); - // The per-layer embedding planes, ONCE per forward, for EVERY layer. - // This is the host touch that is allowed, because it happens before - // anything is dispatched. The tail stays zero, which is what the - // padded rows' PLE input should be. - std::vector plane(static_cast(padded * cfg.ple_dim), 0.0f); - for (std::size_t layer = 0; layer < planes.size(); ++layer) { - std::fill(plane.begin(), plane.end(), 0.0f); - std::copy(planes[layer].begin(), planes[layer].end(), plane.begin()); - api->Check(api->functions().tensor_write(ple_in[layer].get(), - ryzenai_corelib_data_type_fp32, - plane.data(), plane.size(), 0), - "ryzenai_corelib_tensor_write per-layer embedding plane"); - } - - // ---- every window, before the first dispatch - const auto hidden_w = Window(hidden.get(), {padded, cfg.hidden}, 0, "hidden"); - const auto residual_w = Window(residual.get(), {padded, cfg.hidden}, 0, "residual"); - const auto skip_w = Window(skip_sum.get(), {padded, cfg.hidden}, 0, "skip_sum"); - const auto ple_x_w = Window(ple_out.get(), {padded, cfg.hidden}, 0, "ple x"); - std::vector ple_w; - ple_w.reserve(static_cast(cfg.layers)); - for (std::int64_t layer = 0; layer < cfg.layers; ++layer) { - ple_w.push_back(Window(ple_in[static_cast(layer)].get(), - {padded, cfg.ple_dim}, 0, "per-layer embedding")); - } - std::array, 2> geometry; - const auto geometry_for = [&](bool swa) -> const Geometry& { - auto& slot = geometry[swa ? 0 : 1]; - if (!slot) slot = BuildGeometry(swa, padded); - return *slot; - }; - for (std::int64_t layer = 0; layer < cfg.layers; ++layer) - (void)geometry_for(cfg.layer_is_swa[static_cast(layer)]); - // Where V's norm writes, per owning layer and per head: the SAME - // per-head arithmetic as `v_src`, plus `position`. A sharing layer - // gets none -- it projects no V at all, and handing it a window would - // be a way to overwrite somebody else's cache with nothing. - std::vector> v_cache_w( - static_cast(cfg.layers)); - for (std::int64_t layer = 0; layer < cfg.kv_layers; ++layer) { - // THE HEAD SIZE COMES FROM THE LAYER'S OWN GEOMETRY, not from its - // index. The two agree on any model this loads, but they agree by - // construction rather than by definition, and this window is the - // one thing in the pass that would then stride a cache at the - // other geometry's pitch. - const auto head = - geometry_for(cfg.layer_is_swa[static_cast(layer)]).head; - auto& windows = v_cache_w[static_cast(layer)]; - windows.reserve(static_cast(cfg.kv_heads)); - for (std::int64_t kv_head = 0; kv_head < cfg.kv_heads; ++kv_head) { - windows.push_back(Window( - v_cache[static_cast(layer)].get(), {padded, head}, - CacheOffset(kv_head, position, head), "V cache")); - } - } - - // THE ENTIRE DEFINITION OF POISONING IS THIS FLAG, so what sets it is - // load-bearing rather than incidental. IT FLIPS ONLY ON AN ACCEPTED - // CALL -- after `Check` has let the status through -- so it means - // "corelib took a command, and that command is in flight or has run". - // A REFUSED dispatch enqueued nothing: if it is the first of the - // pass there is nothing in flight and the engine is exactly as it - // was, and if it is not the first then the accepted calls before it - // have already set the flag. Setting it BEFORE the `Check` instead -- - // which is what this did until C9's fix round -- makes the refusal of - // the very first dispatch poison an engine nothing has touched, and - // makes `if (submitted)` in the catch indistinguishable from - // `if (true)`, which is how the guard came to have no test. - bool submitted = false; - try { - const auto dispatch = [&](ryzenai_corelib_status status, const char* what, - std::int64_t layer) { - api->Check(status, std::string(what) + " layer " + std::to_string(layer)); - submitted = true; - }; - // LAYER 0'S ATTN_NORM, THE ONE STANDALONE RMSNORM IN THE MODEL. - // Every other input norm is the previous layer's `ple` second - // output; this one reads the embedding with no block before it. - api->Check(api->functions().rmsnorm(stream.get(), residual_w.get(), - input_norm.get(), hidden_w.get()), - "ryzenai_corelib_rmsnorm_bf16 layer 0 input norm"); - submitted = true; - - for (std::int64_t index = 0; index < cfg.layers; ++index) { - const auto slot = static_cast(index); - const auto& weights = layers[slot]; - const bool swa = cfg.layer_is_swa[slot]; - const bool owns = index < cfg.kv_layers; - const auto& kind = geometry_for(swa); - - dispatch(api->functions().matmul(stream.get(), hidden_w.get(), - weights.q.get(), kind.q.get()), - "ryzenai_corelib_matmul_bf16 query", index); - // K AND V ONLY ON A LAYER THAT OWNS A CACHE. A sharing layer - // attends over a cache an earlier layer of its own kind wrote, - // and its descriptor is a kvshare artifact that takes no K - // operand at all. - if (owns) { - dispatch(api->functions().matmul(stream.get(), hidden_w.get(), - weights.k.get(), kind.k.get()), - "ryzenai_corelib_matmul_bf16 key", index); - dispatch(api->functions().matmul(stream.get(), hidden_w.get(), - weights.v.get(), kind.v_dest.get()), - "ryzenai_corelib_matmul_bf16 value", index); - dispatch(api->functions().rmsnorm(stream.get(), kind.k_heads_in.get(), - weights.k_norm.get(), - kind.k_heads_out.get()), - "ryzenai_corelib_rmsnorm_bf16 K-norm", index); - // V'S NORM IS WHAT WRITES THE CACHE, once per KV head, - // because each head's plane is a separate contiguous run - // in both buffers. Its gamma is ALL ONES -- the file - // carries no learned weight -- and dropping the step - // because of that leaves a whole 1/sqrt(mean(v^2)) out of - // every attended value. - for (std::int64_t kv_head = 0; kv_head < cfg.kv_heads; ++kv_head) { - dispatch(api->functions().rmsnorm( - stream.get(), - kind.v_src[static_cast(kv_head)].get(), - weights.v_norm.get(), - v_cache_w[slot][static_cast(kv_head)].get()), - "ryzenai_corelib_rmsnorm_bf16 V-norm", index); - } - } - // QK-NORM, BEFORE ATTENTION, because flat_mha applies the - // rotary itself and Gemma normalizes ahead of it. THESE GAMMAS - // ARE THE ATTENTION SCALE -- constant-valued vectors whose - // product is 0.125 on a sliding layer and 0.0625 on a full one - // -- which is why the descriptor's scale is genuinely 1.0 and - // why skipping a norm whose gamma "looks constant" is an 8x or - // 16x logit error. NOTHING HERE PRE-SCALES Q. - dispatch(api->functions().rmsnorm(stream.get(), kind.q_heads_in.get(), - weights.q_norm.get(), - kind.q_heads_out.get()), - "ryzenai_corelib_rmsnorm_bf16 Q-norm", index); - // TWO ROTARY PAIRS, AND `swa` SELECTS. One pair used - // everywhere still converges to fluent text. - const auto owner = static_cast(cfg.layer_cache_owner[slot]); - void* cos = (swa ? cos_sliding : cos_full).get(); - void* sin = (swa ? sin_sliding : sin_full).get(); - if (plan.paged()) { - dispatch(api->functions().paged_mha( - stream.get(), &plan.paged_attention_desc(swa, /*shared=*/!owns), - kind.q_attn.get(), kind.k_attn.get(), position, cos, sin, - k_cache[owner].get(), v_cache[owner].get(), kind.attn.get()), - "ryzenai_corelib_paged_mha_bf16", index); - } else { - dispatch(api->functions().flat_mha( - stream.get(), - // SELECT, NEVER MUTATE. The plan holds four - // immutable descriptors keyed (geometry, cache - // role); there is deliberately no one-argument - // accessor, so a loop that forgot the cache role - // would not compile. - &plan.attention_desc(swa, /*shared=*/!owns), - kind.q_attn.get(), kind.k_attn.get(), position, cos, sin, - k_cache[owner].get(), v_cache[owner].get(), kind.attn.get()), - "ryzenai_corelib_flat_mha_bf16", index); - } - dispatch(api->functions().matmul(stream.get(), kind.attn.get(), - weights.o.get(), hidden_w.get()), - "ryzenai_corelib_matmul_bf16 output", index); - dispatch(api->functions().rmsnorm(stream.get(), hidden_w.get(), - weights.post_attn_norm.get(), - skip_w.get()), - "ryzenai_corelib_rmsnorm_bf16 post-attention", index); - // THE BLOCK'S OUTPUT IS `skip_sum` (PLANE 0), NOT `normalized` - // (plane 1), and that is the opposite of the silu family. - // Under the gemma post-feedforward topology the trailing norm - // moves INSIDE the residual, DD declares a single output, and - // plane 1 COMES BACK UNTOUCHED -- it is still bound, because - // an unbound plane dispatches a null pointer. Measured with - // the operands the other way round, `normalized` returns a - // plane of EXACT ZEROS, which propagates through `ple` and out - // of all 35 layers as a finite, correctly-shaped vector of - // zeros. `hidden_w` is the plane nothing writes, and `ple` - // overwrites it two lines down. - dispatch(api->functions().ssmlp(stream.get(), skip_w.get(), - residual_w.get(), weights.mlp.get(), - ple_x_w.get(), hidden_w.get()), - "ryzenai_corelib_ssmlp_bf16", index); - // TWO OUTPUTS: the new residual stream, and the NEXT layer's - // normed input. The second is why no layer boundary touches - // the host, and on the LAST layer its gamma is `output_norm` - // rather than a 36th block's -- so `hidden` already carries - // the model's final norm when the loop ends. DO NOT APPLY A - // FINAL NORM AFTERWARDS. - dispatch(api->functions().ple(stream.get(), ple_x_w.get(), - ple_w[slot].get(), weights.ple.get(), - residual_w.get(), hidden_w.get()), - "ryzenai_corelib_ple_bf16", index); - } - - // THE ONLY WAIT IN THE LOOP, and it is here because the next line - // reads on the host. - api->Check(api->functions().stream_synchronize(stream.get()), - "ryzenai_corelib_stream_synchronize layer loop"); - // The LAST live row. Read back as FP32: the tensor is BF16, so - // corelib widens on the way out and there is nothing to unpack. - std::vector last(static_cast(cfg.hidden)); - api->Check(api->functions().tensor_read( - hidden.get(), ryzenai_corelib_data_type_fp32, last.data(), - last.size(), - static_cast((rows - 1) * cfg.hidden)), - "ryzenai_corelib_tensor_read final hidden row"); - RequireFinite(last, "the forward pass"); - - // ---- the head. lm_head runs at M == 1 alone at this model's - // (8, 17) pair -- one row is not a simplification, it is the only - // shape available -- so its input is a separate one-row tensor. - api->Check(api->functions().tensor_write(last_hidden.get(), - ryzenai_corelib_data_type_fp32, - last.data(), last.size(), 0), - "ryzenai_corelib_tensor_write LM head input"); - api->Check(api->functions().matmul(stream.get(), last_hidden.get(), - lm_head.get(), logits.get()), - "ryzenai_corelib_matmul_bf16 LM head"); - api->Check(api->functions().stream_synchronize(stream.get()), - "ryzenai_corelib_stream_synchronize logits"); - std::vector raw(static_cast(cfg.vocab)); - api->Check(api->functions().tensor_read(logits.get(), - ryzenai_corelib_data_type_fp32, - raw.data(), raw.size(), 0), - "ryzenai_corelib_tensor_read logits"); - RequireFinite(raw, "the head's logits"); - // THE SOFTCAP, ON THE HOST, AFTER THE READ-BACK: tanh(x/30)*30, - // with 30 read from the file rather than written here. There is no - // matmul epilogue that does it. Omitting it looks like nothing -- - // the raw logits are finite and well spread, the argmax usually - // agrees, and the text stays fluent -- but the cap changes every - // softmax it feeds. A model that states 0.0 means "no cap", which - // `tanh(x/0)` would answer with a plane of NaNs. - if (cfg.logit_softcap != 0.0f) SoftcapLogits(raw, cfg.logit_softcap); - - buffer out(static_cast(cfg.vocab)); - std::vector as_bf16(raw.size()); - FloatsToBf16(raw, as_bf16); - static_assert(sizeof(bf16) == sizeof(std::uint16_t), - "buffer holds two bytes per logit"); - std::memcpy(out.data(), as_bf16.data(), - as_bf16.size() * sizeof(std::uint16_t)); - position += static_cast(rows); - return out; - } catch (...) { - // WAIT FOR WHAT WAS SUBMITTED BEFORE UNWINDING PAST THE OPERANDS. - // This does not make a mid-chain corelib rejection survivable -- - // nothing does, which is why every check above is before the first - // dispatch -- but a failure raised on the HOST side after a - // successful submission (an allocation, the finite check) must not - // destroy the windows while their commands are in flight. - // - // AND THE ENGINE IS POISONED FROM HERE ON. `submitted` is the - // whole of the condition, and it means an ACCEPTED dispatch: - // a command corelib took, which is in flight or has run, and - // whose effect on the device this function cannot put back. - // - // - a mid-chain corelib rejection is unrecoverable by - // construction (see this function's header: the chain is - // unbarriered, so the unwind passes operands whose commands - // are still in flight). The engine cannot know how much of - // the loop ran, and the KV caches now hold rows for - // positions it will not count. - // - a non-finite result is PERMANENT -- BUT NOT FOR THE REASON - // THIS COMMENT USED TO GIVE. It said "every later token - // would attend over it", and that is not established: - // `flat_mha` is handed `position`, and `position += rows` - // runs only on the success path, so a failed pass's NaN rows - // sit at [position, position + padded) -- outside the window - // any later pass attends over, and a retry at the same - // position overwrites them. The two reasons that DO hold: - // * THE CAUSE IS DETERMINISTIC -- weights, kernel - // selection, shape -- so a retry reproduces it. The flag - // costs a caller nothing that was actually available. - // * `RequireFinite` SEES LESS THAN THE CORRUPTION IT - // DETECTS: the final live row's hidden, and the logits, - // and nothing else. A NaN written into a cache row BELOW - // `position` by an earlier, apparently successful pass is - // exactly what it cannot see -- and `clear_context()` - // resets a POSITION, it does not zero a cache. The - // observable failure understates the damage rather than - // overstating it, which is an argument for poisoning and - // not against it. - // - // Anything refused before a dispatch was ACCEPTED leaves the - // engine exactly as it was and must not poison it: the caller - // errors (an empty request, a token outside the vocabulary, a - // chunked prefill, a request past capacity), plus the refusal of - // the pass's very first dispatch, which enqueued nothing. Turning - // any of those into a mandatory model reload would be a worse bug - // than the one this guards. - // - // THE TWO `tensor_write`s ABOVE ALSO WRITE DEVICE MEMORY, and - // they sit before the `try`. So "a submitted dispatch is the only - // thing here that writes device state" -- which this comment used - // to say -- is false. `residual` and every `ple_in` plane are - // rewritten IN FULL at the head of every pass, so a partial write - // to either is recoverable; the true rule, and the one to apply - // when adding a device write, is "the only thing that writes - // state a later pass reads WITHOUT FIRST REWRITING IT". - // - // BOTH HALVES ARE PINNED, in test_gemma4_engine.cpp: refusing - // dispatch 20 must poison and refusing dispatch 0 must not, and - // a third, structural test keeps the host rejections outside this - // `try` so that the second of those keeps meaning what it says. - // - // `position` and `saved` are cleared for the same reason Phi-4's - // engine clears them: the surviving numbers describe a - // continuation that no longer exists. Every accessor is behind - // `usable()`, so this is defence rather than something a caller - // can observe. - // - // ONE EXCEPTION. A paged prefill chunk past page 0 can come back - // non-finite from an intermittent kernel defect: the chunk - // rewrites only its own page, and `position` has not moved, so - // RunPrompt may run that same chunk again. Unpaged passes and - // page 0 still poison. The message text is RequireFinite's. - bool retry = false; - try { - throw; - } catch (const std::exception& error) { - retry = paged && position > 0 && - std::string(error.what()).find("non-finite") != std::string::npos; - } catch (...) { - } - if (submitted) { - (void)api->functions().stream_synchronize(stream.get()); - if (!retry) { - poisoned = true; - position = 0; - saved.reset(); - } - } - throw; - } - } - - /// \brief one prompt, one page at a time when the cache is paged - /// \note Starts at `position`, which is 0 or a page boundary when `ids` - /// holds more than one row. - buffer RunPrompt(std::span ids) { - if (!paged || ids.size() <= 1 || - (position == 0 && static_cast(ids.size()) <= kPageSize)) - return Run(ids); - buffer out; - const auto step = static_cast(kPageSize); - for (std::size_t start = 0; start < ids.size(); start += step) { - const auto count = std::min(step, ids.size() - start); - const auto chunk = ids.subspan(start, count); - if (position == 0) { - out = Run(chunk); - continue; - } - for (int attempt = 1;; ++attempt) { - try { - out = Run(chunk); - break; - } catch (const std::runtime_error& error) { - const bool again = - std::string(error.what()).find("non-finite") != std::string::npos && - attempt < kPagedChunkAttempts; - if (!again) throw; - std::cout << "[FLM] Gemma 4 paged prefill at position " << position - << " came back non-finite; running it again (" << attempt - << " of " << kPagedChunkAttempts << ")" << std::endl; - } - } - } - return out; - } - - /// \brief refuse a non-finite result, naming where it came from - /// \note CHECKED AT THE PASS'S OWN EXIT AND AGAIN AT THE HEAD, because - /// both are places a caller reads from. Without it an all-NaN result - /// reaches an argmax that answers 0 -- ``, which decodes to the - /// empty string -- for every step, so a CLI prints an empty answer - /// under a token count and a rate and exits 0. - /// \note IT IS THE WEAKER HALF OF THE GUARD AND MUST NOT BE READ AS THE - /// WHOLE OF IT. A first pass can be finite and WRONG -- in range, - /// well spread, correctly bounded, and with a different argmax -- - /// and nothing about such a vector says so. - static void RequireFinite(std::span values, const char* what) { - for (std::size_t i = 0; i < values.size(); ++i) { - if (std::isfinite(values[i])) continue; - throw std::runtime_error(std::string("Gemma 4: ") + what + - " returned a non-finite value at index " + - std::to_string(i)); - } - } - - /// \brief one position's K or V cache row, gathered across the KV heads - buffer ReadCache(bool is_k, int layer, int index) { - usable(); - const auto& cfg = shape; - if (layer < 0 || layer >= cfg.layers || index < 0 || - index >= static_cast(cache_rows)) - throw std::out_of_range("Gemma 4 cache index is out of range"); - // THE CACHE A LAYER READS IS ITS OWNER'S, not its own index. Twenty of - // E2B's layers own none at all. - const auto owner = static_cast( - cfg.layer_cache_owner[static_cast(layer)]); - const auto head = cfg.LayerHeadDim(static_cast(owner)); - auto lease = runtime->AcquireExecution(); - api->Check(api->functions().stream_synchronize(stream.get()), - "ryzenai_corelib_stream_synchronize cache read"); - buffer out(static_cast(cfg.kv_heads * head)); - void* cache = (is_k ? k_cache : v_cache)[owner].get(); - for (std::int64_t kv_head = 0; kv_head < cfg.kv_heads; ++kv_head) { - const auto offset = CacheOffset(kv_head, index, head); - api->Check(api->functions().tensor_read( - cache, ryzenai_corelib_data_type_bf16, - out.data() + kv_head * head, static_cast(head), - offset), - "ryzenai_corelib_tensor_read cache head " + - std::to_string(kv_head)); - } - return out; - } -}; - -gemma4_rai::gemma4_rai(std::unique_ptr impl) : impl_(std::move(impl)) {} - -gemma4_rai::gemma4_rai(std::shared_ptr package, - std::shared_ptr runtime, - const LM_Config& config, std::uint32_t max_length) - : impl_(std::make_unique(std::move(package), std::move(runtime), config, - std::nullopt, max_length)) {} - -std::unique_ptr gemma4_rai::CreateWithShapeForTest( - std::shared_ptr package, - std::shared_ptr runtime, const LM_Config& config, - Gemma4Config shape, std::uint32_t max_length) { - return std::unique_ptr( - new gemma4_rai(std::make_unique(std::move(package), std::move(runtime), - config, std::move(shape), max_length))); -} - -gemma4_rai::~gemma4_rai() = default; - -buffer gemma4_rai::forward(int id) { - auto out = impl_->Run(std::span(&id, 1)); - impl_->history.push_back(id); - return out; -} - -// ONE PATH, NOT TWO, and that is the reference driver's shape rather than -// Phi-4's: `forward(hidden, rows, position, ple_in)` is a single function -// there, and the only thing that distinguishes a prefill from a decode step is -// the row count it is handed. The rules that would otherwise be -// "prefill-only" -- a multi-row pass must start at position 0, and every -// buffer is windowed at one padded extent -- are stated once, in Run, as -// properties of the row count and the position. -buffer gemma4_rai::prefill(std::vector& ids, void*) { - auto& impl = *impl_; - if (impl.position == 0 || ids.size() <= 1) { - if (impl.position == 0) impl.history.clear(); - auto out = impl.RunPrompt(std::span(ids.data(), ids.size())); - impl.history.insert(impl.history.end(), ids.begin(), ids.end()); - return out; - } - // A multi-row pass starts at position 0 or, on a paged cache, on a page - // boundary. So a continuation re-runs the rows since the nearest such - // start ahead of its own: the whole conversation unpaged, and on a paged - // cache only the current page's rows, none at all on a boundary. - impl.usable(); - if (impl.history.size() != static_cast(impl.position)) - throw std::invalid_argument("Gemma 4 cannot continue from position " + - std::to_string(impl.position) + - " without the tokens behind it"); - const int resume = impl.position; - const int from = impl.paged ? resume - resume % static_cast(kPageSize) : 0; - std::vector tokens(impl.history.begin() + from, impl.history.end()); - tokens.insert(tokens.end(), ids.begin(), ids.end()); - impl.position = from; - try { - auto out = impl.RunPrompt(std::span(tokens.data(), tokens.size())); - impl.history.insert(impl.history.end(), ids.begin(), ids.end()); - return out; - } catch (...) { - if (!impl.poisoned) impl.position = resume; - throw; - } -} - -void gemma4_rai::set_context_length(int length) { - impl_->usable(); - if (length < 0 || static_cast(length) > impl_->max_length) - throw std::out_of_range("Gemma 4 context length is out of range"); - impl_->position = length; - if (impl_->history.size() > static_cast(length)) - impl_->history.resize(static_cast(length)); -} - -// An ABI shim, not a capability. load_weights is pure virtual in causal_lm.hpp, -// which is frozen because the engine libraries in src/lib/ are -// prebuilt against it -- adding, removing or reordering a virtual there -// corrupts dispatch at runtime with no compiler error. Nothing calls this: -// FlmBackend loads weights through the concrete engine type, and this engine's -// weights come from the GGUF package it was constructed with. See -// AutoModel/model_backend.hpp. -void gemma4_rai::load_weights(Q4NX&) { - impl_->usable(); - throw std::runtime_error("Gemma 4 rai weights are loaded only from GGUF"); -} - -void gemma4_rai::update_max_length(std::uint32_t max_length) { - impl_->usable(); - if (!max_length || max_length > impl_->cache_rows || - max_length < static_cast(impl_->position)) - throw std::out_of_range("Gemma 4 maximum length is invalid"); - impl_->max_length = max_length; -} - -void gemma4_rai::clear_context() { - impl_->usable(); - impl_->position = 0; - impl_->saved.reset(); - impl_->history.clear(); -} - -buffer gemma4_rai::get_k_cache(int layer, int index) { - return impl_->ReadCache(/*is_k=*/true, layer, index); -} - -buffer gemma4_rai::get_v_cache(int layer, int index) { - return impl_->ReadCache(/*is_k=*/false, layer, index); -} - -int gemma4_rai::get_current_context_length() { - impl_->usable(); - return impl_->position; -} - -int gemma4_rai::checkpoint() { - impl_->usable(); - impl_->saved = impl_->position; - return impl_->position; -} - -int gemma4_rai::restore() { - impl_->usable(); - if (!impl_->saved) return -1; - impl_->position = *impl_->saved; - if (impl_->history.size() > static_cast(impl_->position)) - impl_->history.resize(static_cast(impl_->position)); - return impl_->position; -} - -bool gemma4_rai::poisoned() const noexcept { return impl_ && impl_->poisoned; } - -bool gemma4_rai::loaded_from_cache() const noexcept { - return impl_ && impl_->from_cache; -} - -std::size_t gemma4_rai::weight_slot_count() const noexcept { - return impl_ ? impl_->slot_count : 0; -} - -std::vector gemma4_rai::WeightPlacementsForTest() const { - std::vector placements; - if (!impl_) return placements; - placements.reserve(impl_->slot_count); - // NAME EACH MEMBER SEPARATELY, IN FULL. See the header: this function's - // only value is that it is a SECOND statement of the tensor-to-member - // mapping, written here rather than read off the slot table that - // performed the assignment. A loop over `slots` would agree with the - // assignment by construction and could never disagree with it. - placements.push_back({"blk.0.attn_norm.weight", impl_->input_norm.get()}); - for (std::size_t index = 0; index < impl_->layers.size(); ++index) { - const auto layer = static_cast(index); - const auto& weights = impl_->layers[index]; - placements.push_back({Blk(layer, ".attn_q.weight"), weights.q.get()}); - if (weights.k) placements.push_back({Blk(layer, ".attn_k.weight"), weights.k.get()}); - if (weights.v) placements.push_back({Blk(layer, ".attn_v.weight"), weights.v.get()}); - placements.push_back({Blk(layer, ".attn_q_norm.weight"), weights.q_norm.get()}); - if (weights.k_norm) - placements.push_back({Blk(layer, ".attn_k_norm.weight"), weights.k_norm.get()}); - if (weights.v_norm) - placements.push_back({Blk(layer, ".attn_v_norm(ones)"), weights.v_norm.get()}); - placements.push_back({Blk(layer, ".attn_output.weight"), weights.o.get()}); - placements.push_back( - {Blk(layer, ".post_attention_norm.weight"), weights.post_attn_norm.get()}); - placements.push_back({Blk(layer, ".ssmlp"), weights.mlp.get()}); - placements.push_back({Blk(layer, ".ple"), weights.ple.get()}); - } - placements.push_back({"token_embd.weight", impl_->lm_head.get()}); - return placements; -} - -} // namespace flm::gemma4 diff --git a/src/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp b/src/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp index 7d90cfc50..47ef9e51c 100644 --- a/src/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp +++ b/src/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp @@ -1,246 +1,24 @@ /// \file gemma4_rai_backend.cpp -/// \brief The ryzenai-corelib backend for Gemma 4 (E2B and E4B) +/// \brief Gemma 4 E2B and E4B on corelib's "gemma4" model #include "models/gemma4/rai/aie_next/gemma4_rai_backend.hpp" -#include "models/gemma4/rai/aie_next/gemma4_rai.hpp" -#include "models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp" -#include "rai/corelib_runtime.hpp" -#include "utils/file_access.hpp" -#include "utils/utils.hpp" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include +#include "rai/corelib_model_backend.hpp" namespace flm::gemma4 { -namespace { - -/// \brief read a JSON file, recording the open for the file-access audit -/// \param path the file -/// \return the parsed document -/// \throws std::runtime_error naming the path, when it cannot be opened or -/// parsed -- a package problem has to name the file it is about -nlohmann::json ReadJson(const std::filesystem::path& path) { - flm::file_access::ObserveOpen(path); - std::ifstream input(path, std::ios::binary); - if (!input) throw std::runtime_error("Cannot open " + path.string()); - try { - return nlohmann::json::parse(input); - } catch (const std::exception& error) { - throw std::runtime_error("Cannot parse " + path.string() + ": " + - error.what()); - } -} - -/// \brief case-insensitive ".gguf" test -bool NamesAGguf(std::string_view name) { - constexpr std::string_view suffix = ".gguf"; - if (name.size() <= suffix.size()) return false; - auto tail = name.substr(name.size() - suffix.size()); - return std::equal(tail.begin(), tail.end(), suffix.begin(), - [](unsigned char a, unsigned char b) { - return std::tolower(a) == std::tolower(b); - }); -} - -/// \brief the GGUF this package holds -/// -/// NOT A CONSTANT, and that is the one structural difference from Phi-4's -/// backend. Phi-4 can name `Phi-4-mini-instruct.Q8_0.gguf` outright because -/// one file serves the whole family; Gemma 4's two shipped rows are -/// `gemma-4-E2B-it-Q8_0.gguf` and `gemma-4-E4B-it-Q8_0.gguf` under the SAME -/// family `gemma4e`, so a literal here would load exactly one of them and -/// report "file not found" for the other -- an error naming a file the user -/// never asked for. -/// -/// The catalog entry is the authority: `files` is written by the same -/// model_list.json entry the downloader fetched the package with, so it names -/// the file that is actually on disk. A hand-assembled directory has no -/// catalog entry to consult, so a directory scan is the fallback -- one GGUF -/// in a directory is unambiguous, and anything else is named rather than -/// guessed at. -/// -/// \param model_info the catalog entry, as AutoModel passes it down -/// \param root the model directory -/// \return the path to open -/// \throws std::runtime_error naming what it found, when the answer is not -/// exactly one file -std::filesystem::path RequireGgufPath(const nlohmann::ordered_json& model_info, - const std::filesystem::path& root) { - std::vector named; - if (const auto files = model_info.find("files"); - files != model_info.end() && files->is_array()) { - for (const auto& entry : *files) { - if (entry.is_string() && NamesAGguf(entry.get())) { - named.push_back(entry.get()); - } - } - } - if (named.size() == 1) { - auto path = root / named.front(); - if (!std::filesystem::exists(path)) { - throw std::runtime_error( - "Gemma 4 rai: the catalog entry names " + named.front() + - " but " + path.string() + " does not exist"); - } - return path; - } - if (named.size() > 1) { - std::string list; - for (const auto& name : named) list += (list.empty() ? "" : ", ") + name; - throw std::runtime_error( - "Gemma 4 rai: the catalog entry names " + - std::to_string(named.size()) + " GGUFs (" + list + - "); it must name exactly one"); - } - - // No catalog entry to go on: scan. - std::vector found; - std::error_code failed; - for (const auto& entry : std::filesystem::directory_iterator(root, failed)) { - if (entry.is_regular_file(failed) && - NamesAGguf(entry.path().filename().string())) { - found.push_back(entry.path()); - } - } - if (found.size() == 1) return found.front(); - throw std::runtime_error( - "Gemma 4 rai: the catalog entry names no GGUF and " + root.string() + - " holds " + std::to_string(found.size()) + - " of them; it must hold exactly one"); -} - -/// \brief Gemma 4 on ryzenai-corelib, GGUF weights -class RaiBackend final : public flm::backend::ModelBackend { -public: - explicit RaiBackend(const flm::backend::BackendContext& context) { - // 1. Refuse what this backend cannot do, FIRST -- before a file is - // opened or a runtime acquired. Both are also BackendTraits, so - // AutoModel rejects them earlier still; these are the second line, - // for a caller that reaches the factory directly. - if (!context.config) { - throw std::runtime_error("Gemma 4 rai backend needs an LM_Config"); - } - if (context.enable_preemption) { - throw std::invalid_argument( - "Gemma 4 rai does not support preemption"); - } - if (!ContextFits(context.context_length)) { - throw std::out_of_range( - "Gemma 4 rai context length must be in 1..4096, or a multiple of 4096 up to " + - std::to_string(kMaxPagedContext)); - } - context_length_ = context.context_length; - - // 2-3. Read and cross-validate EVERY source of truth before the - // runtime is acquired or the engine built, so a mismatched - // package fails at load with an exact diagnostic and nothing - // allocated -- never mid-generation. - const std::filesystem::path root(context.model_path); - const auto config = ReadJson(root / "config.json"); - const auto tokenizer_json = ReadJson(root / "tokenizer.json"); - // Phi-4's backend takes tokenizer_config.json from the frontend, - // which has already parsed it. Gemma4e cannot do that: its - // load_model calls _shared_load_backend BEFORE setup_tokenizer, so - // there is nothing parsed to hand down at the moment the factory - // runs. Read it here when it is not supplied, and use the frontend's - // copy when it is, so the file is opened once per load either way. - const auto tokenizer_config = - context.tokenizer_config ? *context.tokenizer_config - : ReadJson(root / "tokenizer_config.json"); - auto package = - Gemma4GgufPackage::Open(RequireGgufPath(context.model_info, root)); - package->ValidateGemma4Contract(config, tokenizer_json, - tokenizer_config); - - // Both ids come from the GGUF, cross-checked by the contract above. - // Generation stops on the GGUF's eos (106, ) and on every other - // id config.json lists (1, ). - forced_eos_ids_.clear(); - for (const auto id : ConfigEosIds(config, package->Config().eos_token_id)) - forced_eos_ids_.push_back(static_cast(id)); - forced_bos_id_ = package->Config().bos_token_id; - - // 4-6. Only now: the runtime, the engine, and a cleared context. - runtime_ = corelib::CorelibRuntime::GetOrCreate( - utils::get_executable_directory()); - auto engine = std::make_unique( - std::move(package), runtime_, *context.config, - context.context_length); - engine->clear_context(); - engine_ = std::move(engine); - } - - ~RaiBackend() override { - // The engine holds its own reference to the runtime, but destroy it - // here anyway: corelib objects must not outlive the API they came - // from, and an engine released after its runtime is the case that - // access-violates. - // - // NOT AN ENUMERATED TEARDOWN, and the difference matters. A - // destructor BODY runs before member destruction, so a list like - // "engine_.reset(); runtime_.reset();" would release the runtime - // ahead of any member the list forgot. What actually holds the order - // is the DECLARATION ORDER below -- runtime_ first, so it is - // destroyed last. See task C9. - engine_.reset(); - } - - causal_lm& engine() override { return *engine_; } - - std::string id() const override { return flm::backend::kRaiBackendId; } - - std::string detail() const override { - return runtime_ ? runtime_->loaded_library_path().string() - : std::string(); - } - - std::uint32_t max_decode_length() const override { - return context_length_ > 0 ? context_length_ - 1 : 0; - } - - bool supports_preemption() const override { return false; } - - /// \note corelib rejects a decode past its own limit, so the extra - /// forward() the FastFlowLM engines want after an EOS token would - /// fail here. - bool forwards_past_eos() const override { return false; } - - bool poisoned() const noexcept override { - return engine_ && engine_->poisoned(); - } - - std::optional> forced_eos_ids() const override { - return forced_eos_ids_; - } - - std::optional forced_bos_id() const override { return forced_bos_id_; } - -private: - /// \brief read from the package at construction; see the constructor - std::uint32_t context_length_{}; - std::vector forced_eos_ids_; - /// \brief read from the package at construction; see the constructor - int forced_bos_id_{}; - // Declared before the engine so it is destroyed AFTER it. - std::shared_ptr runtime_; - std::unique_ptr engine_; -}; - -} // namespace +/// \note A multi-row pass starts at 0 or, past 4096, on a page boundary, so a +/// later prompt re-runs its own page's earlier tokens with it. +/// \note The package states its stop and start tokens by name: and +/// in tokenizer_config.json, and config.json's eos_token_id [1, 106] +/// adds . flm::backend::BackendFactory rai_factory() { - return [](const flm::backend::BackendContext& context) - -> std::unique_ptr { - return std::make_unique(context); - }; + return rai::Factory({ + .model = {"gemma4", "Gemma 4", kRaiContextLimit, rai::Continuation::kFromPage, + rai::Rewind::kSeek}, + .context_fits = ContextFits, + .context_rule = "1..4096, or a multiple of 4096 up to 16384", + .read_bos = true, + }); } } // namespace flm::gemma4 diff --git a/src/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp b/src/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp deleted file mode 100644 index f73390736..000000000 --- a/src/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp +++ /dev/null @@ -1,503 +0,0 @@ -#include "models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp" - -#include "models/gemma4/rai/aie_next/gemma4_rai_constants.hpp" -#include "rai/gguf_file.hpp" - -#include -#include -#include -#include -#include - -namespace flm::gemma4 { -namespace { - -// TWO ROWS, AND THE SECOND IS THE POINT OF THE FIRST. Everything that -// differs between E2B and E4B -- 42 layers against 35, hidden 2560 against -// 1536, two KV heads against one, a uniform FFN width against one that -// steps, a full-attention period of 6 against 5, a different pair of shared -// caches -- is READ FROM THE FILE by later tasks, not stated here. A row -// carries only what no GGUF key states: the vocabulary size, and the three -// group sizes DynamicDispatch's ELF set ships at PDI pair (8, 17). Those -// three were read off the ELF set itself -- not guessed from either model's -// shape -- which is exactly why an unrecognized shape below is rejected -// rather than assigned a plausible-looking row of its own. -constexpr std::array kRows{{ - {"E2B", 35, 1536, 8, 1, 256, 512, kVocabularySize, kGroupSize, kGroupSize, kGroupSize}, - {"E4B", 42, 2560, 8, 2, 256, 512, kVocabularySize, kGroupSize, kGroupSize, kGroupSize}, -}}; - -} // namespace - -std::span Gemma4Rows() { return kRows; } - -const Gemma4Row& RequireRowFor(std::int64_t layers, std::int64_t hidden, - std::int64_t q_heads, std::int64_t kv_heads, - std::int64_t head_dim, - std::int64_t global_head_dim) { - for (const auto& row : kRows) { - if (row.layers == layers && row.hidden == hidden && - row.q_heads == q_heads && row.kv_heads == kv_heads && - row.head_dim == head_dim && row.global_head_dim == global_head_dim) { - return row; - } - } - // Name every shape field the caller asked about, so whoever hits this - // can tell at a glance whether it is a typo or a genuinely new size -- - // and, if the latter, exactly what to go measure before adding a row. - throw std::runtime_error( - "no Gemma 4 row matches layers=" + std::to_string(layers) + - " hidden=" + std::to_string(hidden) + - " q_heads=" + std::to_string(q_heads) + - " kv_heads=" + std::to_string(kv_heads) + - " head_dim=" + std::to_string(head_dim) + - " global_head_dim=" + std::to_string(global_head_dim) + - "; add a row once its group sizes have been read off the ELF set at " - "PDI pair (8, 17). Guessing them would replace this message with a " - "dispatch failure far from here."); -} - -std::vector DeriveFfnWidths(const flm::rai::GgufFile& file, - std::int64_t layers) { - std::vector widths; - widths.reserve(static_cast(layers)); - for (std::int64_t layer = 0; layer < layers; ++layer) { - const auto name = "blk." + std::to_string(layer) + ".ffn_gate.weight"; - const auto shape = file.Tensor(name).shape; - if (shape.size() != 2) { - throw std::runtime_error( - name + " is not 2-D; Gemma 4's MLP width is read from this " - "tensor's own shape, never from gemma4.feed_forward_length"); - } - // shape[0], NOT shape[1]. `GgufFile::ShapeOf` REVERSES the GGUF's - // own dims, and the file stores them fastest-varying (input) first, - // so a real file's `ShapeOf` is [OUT, IN]: `blk.0.ffn_gate.weight` - // reads [1536, 6144] in the file and [6144, 1536] here. Reading - // shape[1] returns `hidden` on every layer of both models -- a - // plausible, uniform, wrong width, which is precisely what the - // transposed fixture made look right. - widths.push_back(shape[0]); - } - return widths; -} - -std::vector DeriveLayerIsSwa(const flm::rai::GgufFile& file, - std::int64_t layers, std::int64_t q_heads, - std::int64_t head_dim, - std::int64_t global_head_dim) { - // NOT a period. E2B's full-attention layers recur every 5th layer, - // E4B's every 6th -- a pattern generated from one would fit the other's - // weights nowhere. The Q projection's own width says which geometry a - // layer is, and that width is the same on both models even though their - // hidden sizes differ (see the header doc comment). - const std::int64_t sliding_width = q_heads * head_dim; - const std::int64_t full_width = q_heads * global_head_dim; - std::vector is_swa; - is_swa.reserve(static_cast(layers)); - for (std::int64_t layer = 0; layer < layers; ++layer) { - const auto name = "blk." + std::to_string(layer) + ".attn_q.weight"; - const auto shape = file.Tensor(name).shape; - if (shape.size() != 2) { - throw std::runtime_error( - name + " is not 2-D; Gemma 4's attention geometry (sliding " - "vs full-attention) is read from this tensor's own width, " - "never from a derived period"); - } - // shape[0] is the OUTPUT width -- see DeriveFfnWidths above for why. - // This one at least fails loudly when read the other way round - // (`hidden` matches neither geometry), which is how the transposition - // was finally caught; DeriveFfnWidths just returns a wrong number. - if (shape[0] == sliding_width) { - is_swa.push_back(true); - } else if (shape[0] == full_width) { - is_swa.push_back(false); - } else { - throw std::runtime_error( - name + " has width " + std::to_string(shape[0]) + - ", which is neither the sliding geometry (" + - std::to_string(sliding_width) + ") nor the full-attention one (" + - std::to_string(full_width) + ")"); - } - } - return is_swa; -} - -std::vector DeriveCacheOwners(const std::vector& layer_is_swa, - std::int64_t kv_layers) { - const auto layers = static_cast(layer_is_swa.size()); - if (kv_layers <= 0 || kv_layers > layers) { - throw std::runtime_error( - "kv_layers " + std::to_string(kv_layers) + " is outside 1.." + - std::to_string(layers) + " -- this is the OWNING count " - "(block_count - shared_kv_layers), not the raw " - "gemma4.attention.shared_kv_layers metadata value"); - } - // The last own-cache layer of each kind, tracked while scanning only the - // owning prefix -- never a period, for the same reason DeriveLayerIsSwa - // does not use one: the two shipped models disagree on it. - std::int64_t last_sliding = -1; - std::int64_t last_full = -1; - std::vector owners(static_cast(layers), -1); - for (std::int64_t layer = 0; layer < kv_layers; ++layer) { - owners[static_cast(layer)] = layer; - (layer_is_swa[static_cast(layer)] ? last_sliding : last_full) = layer; - } - for (std::int64_t layer = kv_layers; layer < layers; ++layer) { - const bool swa = layer_is_swa[static_cast(layer)]; - const std::int64_t owner = swa ? last_sliding : last_full; - if (owner < 0) { - throw std::runtime_error( - "layer " + std::to_string(layer) + " shares a " + - (swa ? "sliding" : "full-attention") + - " cache, but no layer among the first " + - std::to_string(kv_layers) + " owns one of that kind"); - } - owners[static_cast(layer)] = owner; - } - return owners; -} - -std::int64_t Gemma4Config::LayerHeadDim(std::int64_t layer) const { - return layer_is_swa.at(static_cast(layer)) ? head_dim : global_head_dim; -} - -namespace { - -[[noreturn]] void Fail(std::string_view field, std::string actual, std::string expected) { - throw std::runtime_error(std::string(field) + ": actual " + actual + - ", expected " + expected); -} - -std::string JsonText(const nlohmann::json& value) { return value.dump(); } - -/// \note Not "== \"gemma4\"" on the nose: `config.json`'s own architecture -/// field is an `architectures` array of HF class names (e.g. -/// "Gemma4ForConditionalGeneration"), not the GGUF's short -/// `general.architecture` string. A `model_type` field, if present, -/// is checked directly instead. -bool ArchitectureNamesGemma4(const nlohmann::json& config) { - const auto names_gemma4 = [](const std::string& text) { - return text.find("Gemma4") != std::string::npos; - }; - if (const auto it = config.find("architectures"); it != config.end() && it->is_array()) { - for (const auto& entry : *it) { - if (entry.is_string() && names_gemma4(entry.get())) return true; - } - } - if (const auto it = config.find("model_type"); it != config.end() && it->is_string()) { - const auto model_type = it->get(); - return model_type == "gemma4" || names_gemma4(model_type); - } - return false; -} - -} // namespace - -std::vector ConfigEosIds(const nlohmann::json& config, std::int32_t gguf_eos) { - std::vector ids{gguf_eos}; - const auto it = config.find("eos_token_id"); - if (it != config.end() && it->is_array()) { - for (const auto& id : *it) { - if (id.is_number_integer() && - std::find(ids.begin(), ids.end(), id.get()) == ids.end()) - ids.push_back(id.get()); - } - } - return ids; -} - -namespace { - -void RequireJsonUnsigned(const nlohmann::json& object, std::string_view key, - std::uint64_t expected) { - const auto it = object.find(std::string(key)); - if (it == object.end()) Fail(key, "missing", std::to_string(expected)); - std::uint64_t actual; - if (it->is_number_unsigned()) { - actual = it->get(); - } else if (it->is_number_integer()) { - const auto signed_value = it->get(); - if (signed_value < 0) - Fail(key, JsonText(*it), "non-negative integer " + std::to_string(expected)); - actual = static_cast(signed_value); - } else { - Fail(key, JsonText(*it), "integer " + std::to_string(expected)); - } - if (actual != expected) Fail(key, std::to_string(actual), std::to_string(expected)); -} - -/// \brief where the fields describing the TEXT model live -/// -/// The real `config.json` is a MULTIMODAL config: `architectures`, -/// `model_type` and `eos_token_id` sit at the top level, but every field -/// that describes the text model -- `num_hidden_layers`, `vocab_size`, -/// `intermediate_size`, ... -- is nested one level down, under -/// `"text_config"` (confirmed against the real file; see -/// verified-gguf-facts.md's "config.json is MULTIMODAL" section). An -/// earlier version of this check read `num_hidden_layers` off the top -/// level directly and rejected every real Gemma 4 package as a result -- -/// exactly the failure mode this contract exists to prevent, inverted. -/// -/// Falls back to the top level itself when `text_config` is absent or not -/// an object, in case a hypothetical flat config ever ships. Every real -/// file measured so far has `text_config`, so that fallback is untested -/// insurance, not the expected path -- it exists only so a flat config does -/// not regress into "missing text_config" instead of "missing -/// num_hidden_layers". -const nlohmann::json& TextModelConfig(const nlohmann::json& config) { - if (const auto it = config.find("text_config"); it != config.end() && it->is_object()) - return *it; - return config; -} - -} // namespace - -Gemma4Config ConfigFromMetadata(const flm::rai::GgufFile& file) { - const auto architecture = file.String("general.architecture"); - if (architecture != "gemma4") { - throw std::runtime_error("general.architecture: actual '" + architecture + - "', expected 'gemma4'"); - } - - const auto layers = static_cast(file.Unsigned("gemma4.block_count")); - const auto hidden = static_cast(file.Unsigned("gemma4.embedding_length")); - const auto q_heads = - static_cast(file.Unsigned("gemma4.attention.head_count")); - const auto kv_heads = - static_cast(file.Unsigned("gemma4.attention.head_count_kv")); - - // The GGUF's names run OPPOSITE to ours: key_length is the FULL-attention - // layers' head size and key_length_swa the SLIDING one. Worth knowing at - // this one place and nowhere else -- reading them the obvious way round - // swaps 256 and 512 through the whole model, silently. - const auto global_head_dim = - static_cast(file.Unsigned("gemma4.attention.key_length")); - const auto head_dim = - static_cast(file.Unsigned("gemma4.attention.key_length_swa")); - - // gemma4.attention.shared_kv_layers counts the layers that SHARE a - // cache, NOT the layers that OWN one -- the complement of what - // DeriveCacheOwners's kv_layers parameter means. This subtraction is the - // one correction that matters most in this whole file: taking the raw - // key would give 20 on E2B and 18 on E4B where 15 and 24 are correct, - // and every shared layer would silently bind the wrong KV cache. See - // DeriveCacheOwners's doc comment and verified-gguf-facts.md. - const auto shared_kv_layers = - static_cast(file.Unsigned("gemma4.attention.shared_kv_layers")); - const auto kv_layers = layers - shared_kv_layers; - - const auto sliding_window = - static_cast(file.Unsigned("gemma4.attention.sliding_window")); - // A file that does not state its context gets one page, as corelib's - // driver does. - const auto context_length = file.HasMetadata("gemma4.context_length") - ? static_cast(file.Unsigned("gemma4.context_length")) - : kMaxSequenceLength; - const auto eps = - static_cast(file.Number("gemma4.attention.layer_norm_rms_epsilon")); - const auto rope_theta = file.Number("gemma4.rope.freq_base"); - const auto rope_theta_swa = file.Number("gemma4.rope.freq_base_swa"); - const auto logit_softcap = static_cast(file.Number("gemma4.final_logit_softcapping")); - const auto eos_token_id = - static_cast(file.Unsigned("tokenizer.ggml.eos_token_id")); - const auto bos_token_id = - static_cast(file.Unsigned("tokenizer.ggml.bos_token_id")); - const auto add_bos = file.Boolean("tokenizer.ggml.add_bos_token"); - - auto layer_intermediate = DeriveFfnWidths(file, layers); - auto layer_is_swa = DeriveLayerIsSwa(file, layers, q_heads, head_dim, global_head_dim); - auto layer_cache_owner = DeriveCacheOwners(layer_is_swa, kv_layers); - - const auto& row = RequireRowFor(layers, hidden, q_heads, kv_heads, head_dim, global_head_dim); - - // Named both numbers: which one is wrong (a missing/extra tensor vs. a - // wrong layer count) is not obvious from either alone. - const auto expected_tensor_count = static_cast(6 + layers * 17); - if (file.TensorCount() != expected_tensor_count) { - throw std::runtime_error( - "tensor count: actual " + std::to_string(file.TensorCount()) + - ", expected " + std::to_string(expected_tensor_count) + " (6 + " + - std::to_string(layers) + " layers * 17)"); - } - - // per_layer_token_embd's FILE dims are (layers * ple_dim, vocab) -- the - // layer axis folded into the fastest-varying extent, not a separate - // dimension. `ShapeOf` reverses those, so the layer axis is shape[1] - // here and shape[0] is the vocabulary. ple_dim is 256 on both shipped - // models: 8960 / 35 on E2B and 10752 / 42 on E4B. - // - // Reading shape[0] instead would divide 262144 by the layer count -- - // 7489.8 on E2B, which is not an integer, so THIS one would at least - // throw; on a hypothetical model whose vocabulary happened to divide, - // it would silently produce a ple_dim off by a factor of a thousand. - const auto ple_shape = file.Tensor("per_layer_token_embd.weight").shape; - if (ple_shape.size() != 2 || layers <= 0 || ple_shape[1] % layers != 0) { - throw std::runtime_error( - "per_layer_token_embd.weight's layer axis (" + - (ple_shape.size() != 2 ? std::string("") : std::to_string(ple_shape[1])) + - ") is not a multiple of the layer count (" + std::to_string(layers) + - "); ple_dim is derived as ShapeOf[1] / layers, ShapeOf being the " - "reverse of the file's own dims"); - } - const auto ple_dim = ple_shape[1] / layers; - - Gemma4Config config{}; - config.row_name = row.name; - config.layers = layers; - config.hidden = hidden; - config.q_heads = q_heads; - config.kv_heads = kv_heads; - config.head_dim = head_dim; - config.global_head_dim = global_head_dim; - config.layer_intermediate = std::move(layer_intermediate); - config.layer_is_swa = std::move(layer_is_swa); - config.layer_cache_owner = std::move(layer_cache_owner); - config.kv_layers = kv_layers; - config.sliding_window = sliding_window; - config.vocab = row.vocab; - config.ple_dim = ple_dim; - config.rope_theta = rope_theta; - config.rope_theta_swa = rope_theta_swa; - config.eps = eps; - config.logit_softcap = logit_softcap; - config.eos_token_id = eos_token_id; - config.bos_token_id = bos_token_id; - config.add_bos = add_bos; - config.group = row.group; - config.head_group = row.head_group; - config.ple_group = row.ple_group; - config.context_length = context_length; - return config; -} - -struct Gemma4GgufPackage::Impl { - std::shared_ptr file; - Gemma4Config config; -}; - -Gemma4GgufPackage::Gemma4GgufPackage(std::unique_ptr impl) : impl_(std::move(impl)) {} -Gemma4GgufPackage::~Gemma4GgufPackage() = default; - -std::shared_ptr Gemma4GgufPackage::Open( - const std::filesystem::path& gguf_path) { - auto impl = std::make_unique(); - impl->file = flm::rai::GgufFile::Open(gguf_path); - impl->config = ConfigFromMetadata(*impl->file); - return std::shared_ptr(new Gemma4GgufPackage(std::move(impl))); -} - -const Gemma4Config& Gemma4GgufPackage::Config() const noexcept { return impl_->config; } -const flm::rai::GgufFile& Gemma4GgufPackage::File() const noexcept { return *impl_->file; } -const std::filesystem::path& Gemma4GgufPackage::Path() const noexcept { - return impl_->file->Path(); -} - -void Gemma4GgufPackage::ValidateGemma4Contract(const nlohmann::json& config, - const nlohmann::json& tokenizer, - const nlohmann::json& tokenizer_config) const { - const auto& cfg = impl_->config; - - // architecture and model_type describe the whole (multimodal) package - // and live at the top level even in the real, nested config.json. - if (!ArchitectureNamesGemma4(config)) { - Fail("config.json architecture", - "neither 'architectures' nor 'model_type' names Gemma 4", - "an entry containing \"Gemma4\" (or model_type == \"gemma4\")"); - } - - // Everything below describes the TEXT model and is read from - // text_config, falling back to the top level -- see TextModelConfig's - // doc comment for why, and verified-gguf-facts.md for the real file - // structure that made this necessary. - const auto& text_config = TextModelConfig(config); - - // num_hidden_layers is the one field below that is FATAL if missing, - // and deliberately so: it is the ONLY cross-check field that can catch - // an E2B/E4B swap at all. vocab_size (262144) and eos_token_id (106) - // are IDENTICAL between the two shipped rows -- a config.json that - // named the wrong row but was copy-pasted from within the same family - // would still agree with the GGUF on both of those. Layer count (35 vs - // 42) is where the two models actually disagree, so this is the field - // this check cannot afford to make optional. A field that is silently - // optional is a check that does not exist. - RequireJsonUnsigned(text_config, "num_hidden_layers", - static_cast(cfg.layers)); - - // Deliberately NOT checked, ever: intermediate_size. It says 6144 for - // E2B, which is true only of layers 0-14 -- layers 15-34 actually step - // to 12288 (see DeriveFfnWidths's doc comment for the same trap in the - // GGUF's own metadata key). Comparing it to anything derived here would - // reject a correct E2B package. - - // vocab_size and eos_token_id are cross-checked WHEN PRESENT and - // silently tolerated when absent -- a deliberate choice, not an - // oversight: since both are identical across every row this port - // supports (see the comment on num_hidden_layers above), neither one - // is capable of catching a mismatched package on its own, so refusing - // to load a package that simply omits them would cost real packages - // for no detection gained in return. Their value is in catching a - // config.json that disagrees with the GGUF for some OTHER reason -- - // hand-edited, corrupted, or copied from an unrelated model -- when it - // does bother to state them. - if (text_config.contains("vocab_size")) - RequireJsonUnsigned(text_config, "vocab_size", static_cast(cfg.vocab)); - // eos_token_id is top-level, not under text_config. Google's revisions - // disagree on its type: 106 in some, [1, 106] in others. An array must - // contain the GGUF's eos. - if (config.contains("eos_token_id") && config.at("eos_token_id").is_array()) { - const auto& ids = config.at("eos_token_id"); - const auto want = static_cast(cfg.eos_token_id); - const bool found = std::any_of(ids.begin(), ids.end(), [&](const auto& id) { - return id.is_number_integer() && id.template get() == want; - }); - if (!found) - Fail("eos_token_id", JsonText(ids), "an array containing " + std::to_string(want)); - } else if (config.contains("eos_token_id")) { - RequireJsonUnsigned(config, "eos_token_id", static_cast(cfg.eos_token_id)); - } - // AND THE BOS ID SITS IN THE OTHER PLACE. It is nested under - // `text_config` in every real file and absent from the top level, which - // is the exact mirror of the eos above -- so `TextModelConfig` is the - // object to ask, and asking `config` would silently check nothing. The - // backend's forced_bos_id() rests on this comparison, because - // tokenizer_config.json states no bos_token_id at all for either shipped - // row and so cannot be the source. - if (text_config.contains("bos_token_id")) - RequireJsonUnsigned(text_config, "bos_token_id", static_cast(cfg.bos_token_id)); - - if (tokenizer.contains("model") && tokenizer.at("model").contains("vocab")) { - const auto& vocab = tokenizer.at("model").at("vocab"); - if (vocab.is_object() && vocab.size() != static_cast(cfg.vocab)) { - Fail("tokenizer.json model.vocab size", std::to_string(vocab.size()), - std::to_string(cfg.vocab)); - } - // The third source for each id: tokenizer_config.json NAMES the - // token and tokenizer.json says which id that name has. Both are - // checked only when the package bothers to state them, for the - // reason given above vocab_size. - const auto check_named_token = [&](const char* key, const char* what, - const std::vector& accepted) { - if (!vocab.is_object()) return; - const auto named = tokenizer_config.find(key); - if (named == tokenizer_config.end() || !named->is_string()) return; - const auto token = named->get(); - const auto token_it = vocab.find(token); - if (token_it == vocab.end()) return; - const auto id = token_it->get(); - if (std::find(accepted.begin(), accepted.end(), id) == accepted.end()) { - std::string expected; - for (const auto a : accepted) - expected += (expected.empty() ? "" : " or ") + std::to_string(a); - Fail(what, std::to_string(id), expected); - } - }; - // Newer revisions name "" (1) here while the GGUF stops on - // "" (106); both are in config.json's eos_token_id array. - check_named_token("eos_token", "tokenizer_config.json eos_token", - ConfigEosIds(config, cfg.eos_token_id)); - check_named_token("bos_token", "tokenizer_config.json bos_token", - {cfg.bos_token_id}); - } -} - -} // namespace flm::gemma4 diff --git a/src/common/models/gemma4/rai/aie_next/gemma4_rai_host.cpp b/src/common/models/gemma4/rai/aie_next/gemma4_rai_host.cpp deleted file mode 100644 index 842dd38b8..000000000 --- a/src/common/models/gemma4/rai/aie_next/gemma4_rai_host.cpp +++ /dev/null @@ -1,501 +0,0 @@ -#include "models/gemma4/rai/aie_next/gemma4_rai_host.hpp" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#if defined(_M_X64) || defined(__x86_64__) -#define FLM_GEMMA4_X86 1 -#include -#if defined(_MSC_VER) -#include -#define FLM_GEMMA4_AVX512 -#else -#include -#define FLM_GEMMA4_AVX512 __attribute__((target("avx512f,avx2"))) -#endif -#endif - -namespace flm::gemma4 { -namespace { - -// IEEE binary16 -> float32. Q8_0 blocks store their scale as one fp16 value; -// this mirrors the bit layout phi4_rai_host.cpp already decodes Q8_0 with, -// so the two model ports agree on what a Q8_0 block means. -float HalfToFloat(std::uint16_t half) { - const std::uint32_t sign = static_cast(half & 0x8000) << 16; - const std::uint32_t exponent = (half >> 10) & 0x1f; - std::uint32_t fraction = half & 0x03ff; - std::uint32_t bits; - if (exponent == 0) { - if (fraction == 0) { - bits = sign; - } else { - int shift = 0; - while ((fraction & 0x0400) == 0) { - fraction <<= 1; - ++shift; - } - fraction &= 0x03ff; - bits = sign | (static_cast(127 - 14 - shift) << 23) | - (fraction << 13); - } - } else if (exponent == 0x1f) { - bits = sign | 0x7f800000 | (fraction << 13); - } else { - bits = sign | ((exponent + (127 - 15)) << 23) | (fraction << 13); - } - return std::bit_cast(bits); -} - -// THE PROJECTION: out[r, c] = sum_i x[r, i] * bf16(w[c, i]) for columns -// [c_begin, c_end), accumulated in DOUBLE. The reference comparison in -// test_gemma4_real_gguf.cpp holds these planes to 1e-6 absolute, and FP32 -// accumulation over hidden = 2560 is itself ~1e-6 -- so the precision stays and -// the speed comes from blocking, SIMD and threads. -constexpr std::size_t kProjectColumnTile = 64; -constexpr std::size_t kProjectRowBlock = 16; - -void ProjectScalar(const float* x, const std::uint16_t* w, float* out, std::size_t rows, - std::size_t c_begin, std::size_t c_end, std::size_t width, - std::size_t hidden) { - for (std::size_t row = 0; row < rows; ++row) { - const float* x_row = x + row * hidden; - for (std::size_t column = c_begin; column < c_end; ++column) { - const std::uint16_t* weights = w + column * hidden; - double sum = 0.0; - for (std::size_t i = 0; i < hidden; ++i) { - // BF16 widens EXACTLY: it is the top 16 bits of a binary32. - const auto bits = static_cast(weights[i]) << 16; - sum += static_cast(x_row[i]) * - static_cast(std::bit_cast(bits)); - } - out[row * width + column] = static_cast(sum); - } - } -} - -#if defined(FLM_GEMMA4_X86) -bool CpuHasAvx512F() { - int regs[4] = {}; -#if defined(_MSC_VER) - __cpuid(regs, 0); - if (regs[0] < 7) return false; - __cpuid(regs, 1); - const bool osxsave = (regs[2] >> 27) & 1; - __cpuidex(regs, 7, 0); - const bool avx512f = (regs[1] >> 16) & 1, avx2 = (regs[1] >> 5) & 1; - if (!osxsave || !avx512f || !avx2) return false; - const auto xcr0 = _xgetbv(0); -#else - unsigned a, b, c, d; - if (!__get_cpuid_count(7, 0, &a, &b, &c, &d)) return false; - const bool avx512f = (b >> 16) & 1, avx2 = (b >> 5) & 1; - if (!__get_cpuid(1, &a, &b, &c, &d)) return false; - const bool osxsave = (c >> 27) & 1; - if (!osxsave || !avx512f || !avx2) return false; - unsigned lo, hi; - __asm__("xgetbv" : "=a"(lo), "=d"(hi) : "c"(0)); - const unsigned long long xcr0 = (static_cast(hi) << 32) | lo; -#endif - (void)regs; - // XMM, YMM, opmask, ZMM_Hi256 and Hi16_ZMM state all enabled by the OS. - return (xcr0 & 0xE6) == 0xE6; -} - -// NR rows x 4 columns, eight doubles per lane. Columns past `columns` repeat -// the last one and are not written. -template -FLM_GEMMA4_AVX512 void ProjectTileAvx512(const double* x, const std::uint16_t* w, - float* out, std::size_t column, - std::size_t columns, std::size_t width, - std::size_t hidden) { - __m512d acc[NR][4]; - for (auto& row : acc) - for (auto& lane : row) lane = _mm512_setzero_pd(); - const std::uint16_t* wp[4]; - for (std::size_t k = 0; k < 4; ++k) - wp[k] = w + (column + std::min(k, columns - 1)) * hidden; - const std::size_t body = hidden & ~std::size_t{7}; - for (std::size_t i = 0; i < body; i += 8) { - __m512d wv[4]; - for (std::size_t k = 0; k < 4; ++k) { - const __m128i h = _mm_loadu_si128(reinterpret_cast(wp[k] + i)); - wv[k] = _mm512_cvtps_pd( - _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(h), 16))); - } - for (int j = 0; j < NR; ++j) { - const __m512d xv = _mm512_loadu_pd(x + static_cast(j) * hidden + i); - for (std::size_t k = 0; k < 4; ++k) acc[j][k] = _mm512_fmadd_pd(xv, wv[k], acc[j][k]); - } - } - for (int j = 0; j < NR; ++j) { - const double* x_row = x + static_cast(j) * hidden; - for (std::size_t k = 0; k < columns; ++k) { - double sum = _mm512_reduce_add_pd(acc[j][k]); - for (std::size_t i = body; i < hidden; ++i) { - const auto bits = static_cast(wp[k][i]) << 16; - sum += x_row[i] * static_cast(std::bit_cast(bits)); - } - out[static_cast(j) * width + column + k] = static_cast(sum); - } - } -} - -FLM_GEMMA4_AVX512 void ProjectAvx512(const float* x, const std::uint16_t* w, float* out, - std::size_t rows, std::size_t c_begin, - std::size_t c_end, std::size_t width, - std::size_t hidden, std::vector& block) { - block.resize(kProjectRowBlock * hidden); - for (std::size_t r0 = 0; r0 < rows; r0 += kProjectRowBlock) { - const std::size_t nr = std::min(kProjectRowBlock, rows - r0); - for (std::size_t i = 0; i < nr * hidden; ++i) block[i] = x[r0 * hidden + i]; - for (std::size_t c = c_begin; c < c_end; c += 4) { - const std::size_t nc = std::min(4, c_end - c); - std::size_t r = 0; - float* o = out + r0 * width; - for (; r + 4 <= nr; r += 4) - ProjectTileAvx512<4>(block.data() + r * hidden, w, o + r * width, c, nc, width, hidden); - const double* xr = block.data() + r * hidden; - switch (nr - r) { - case 3: ProjectTileAvx512<3>(xr, w, o + r * width, c, nc, width, hidden); break; - case 2: ProjectTileAvx512<2>(xr, w, o + r * width, c, nc, width, hidden); break; - case 1: ProjectTileAvx512<1>(xr, w, o + r * width, c, nc, width, hidden); break; - default: break; - } - } - } -} -#endif - -// Column tiles are dealt round-robin to one thread per physical core (the -// weights, 55 MB on E4B, are read once per call either way). -void Project(const float* x, const std::uint16_t* w, float* out, std::size_t rows, - std::size_t width, std::size_t hidden) { -#if defined(FLM_GEMMA4_X86) - static const bool avx512 = CpuHasAvx512F(); -#else - constexpr bool avx512 = false; -#endif - const std::size_t tiles = (width + kProjectColumnTile - 1) / kProjectColumnTile; - const std::size_t cores = std::max(1u, std::thread::hardware_concurrency() / 2); - const std::size_t workers = std::min(tiles, cores); - - std::exception_ptr failure; - std::mutex failure_mutex; - const auto run = [&](std::size_t worker) { - try { - std::vector block; - for (std::size_t tile = worker; tile < tiles; tile += workers) { - const std::size_t c_begin = tile * kProjectColumnTile; - const std::size_t c_end = std::min(width, c_begin + kProjectColumnTile); -#if defined(FLM_GEMMA4_X86) - if (avx512) { - ProjectAvx512(x, w, out, rows, c_begin, c_end, width, hidden, block); - continue; - } -#endif - ProjectScalar(x, w, out, rows, c_begin, c_end, width, hidden); - } - } catch (...) { - std::lock_guard lock(failure_mutex); - if (!failure) failure = std::current_exception(); - } - }; - std::vector pool; - pool.reserve(workers - 1); - try { - for (std::size_t worker = 1; worker < workers; ++worker) pool.emplace_back(run, worker); - } catch (...) { - for (auto& thread : pool) thread.join(); - throw; - } - run(0); - for (auto& thread : pool) thread.join(); - if (failure) std::rethrow_exception(failure); -} - -} // namespace - -RopeTables MakeRopeTables(std::int64_t head_dim, double theta, - std::int64_t positions, - std::span frequency_factors) { - if (head_dim <= 0 || head_dim % 2 != 0) { - throw std::runtime_error("rope head_dim must be positive and even, got " + - std::to_string(head_dim)); - } - const std::int64_t half = head_dim / 2; - // ONE FACTOR PER ROTARY PAIR. The driver checks the same thing and says - // the same thing: "rope frequency factors are (n,), not (half,) -- one per - // rotary pair of a `rope_dim`-wide head". A vector of another length is a - // file this code has not been read against, not a value to stretch. - if (!frequency_factors.empty() && - frequency_factors.size() != static_cast(half)) { - throw std::runtime_error( - "rope frequency factors are " + std::to_string(frequency_factors.size()) + - " long, not " + std::to_string(half) + - " -- one per rotary pair of a " + std::to_string(head_dim) + "-wide head"); - } - RopeTables tables; - tables.head_dim = head_dim; - tables.positions = positions; - tables.cos.resize(static_cast(positions * half)); - tables.sin.resize(static_cast(positions * half)); - for (std::int64_t pos = 0; pos < positions; ++pos) { - for (std::int64_t i = 0; i < half; ++i) { - double freq = 1.0 / std::pow(theta, - (2.0 * static_cast(i)) / static_cast(head_dim)); - // DIVIDES. See the header: this is `theta /= freq_factors[i/2]` in - // llama.cpp and `freq = freq / factors` in the reference driver. - // Multiplying instead produces a table that is finite, bounded by - // 1, indistinguishable by any magnitude check, and wrong from - // position 1 on. - if (!frequency_factors.empty()) - freq /= static_cast(frequency_factors[static_cast(i)]); - const double angle = static_cast(pos) * freq; - const auto index = static_cast(pos * half + i); - tables.cos[index] = static_cast(std::cos(angle)); - tables.sin[index] = static_cast(std::sin(angle)); - } - } - return tables; -} - -void RmsNorm(std::span x, std::span gamma, float epsilon, - std::span out) { - if (gamma.size() != x.size()) { - throw std::runtime_error("RmsNorm gamma width " + std::to_string(gamma.size()) + - " does not match x width " + std::to_string(x.size())); - } - if (out.size() != x.size()) { - throw std::runtime_error("RmsNorm out width " + std::to_string(out.size()) + - " does not match x width " + std::to_string(x.size())); - } - double sum_of_squares = 0.0; - for (const float value : x) sum_of_squares += double(value) * value; - const double mean_square = sum_of_squares / static_cast(x.size()); - const double scale = 1.0 / std::sqrt(mean_square + static_cast(epsilon)); - for (std::size_t i = 0; i < x.size(); ++i) { - out[i] = static_cast(static_cast(x[i]) * scale * - static_cast(gamma[i])); - } -} - -void DecodeQ8Row(std::span row, std::span out) { - constexpr std::size_t kBlockBytes = 34; // 2-byte fp16 scale + 32 int8 codes. - constexpr std::size_t kBlockElements = 32; - if (row.size() % kBlockBytes != 0) { - throw std::runtime_error("Q8_0 row byte length " + std::to_string(row.size()) + - " is not a multiple of " + std::to_string(kBlockBytes)); - } - const std::size_t blocks = row.size() / kBlockBytes; - const std::size_t expected_out = blocks * kBlockElements; - if (out.size() != expected_out) { - throw std::runtime_error("Q8_0 decode out width " + std::to_string(out.size()) + - " does not match expected width " + - std::to_string(expected_out)); - } - for (std::size_t block = 0; block < blocks; ++block) { - const std::byte* encoded = row.data() + block * kBlockBytes; - std::uint16_t scale_bits; - std::memcpy(&scale_bits, encoded, sizeof(scale_bits)); - const float scale = HalfToFloat(scale_bits); - for (std::size_t element = 0; element < kBlockElements; ++element) { - const auto code = static_cast( - std::to_integer(encoded[2 + element])); - out[block * kBlockElements + element] = scale * static_cast(code); - } - } -} - -void FloatsToBf16(std::span values, std::span out) { - if (out.size() != values.size()) { - throw std::runtime_error("FloatsToBf16 out width " + std::to_string(out.size()) + - " does not match values width " + - std::to_string(values.size())); - } - for (std::size_t i = 0; i < values.size(); ++i) { - std::uint32_t bits = std::bit_cast(values[i]); - bits += 0x7FFF + ((bits >> 16) & 1); - out[i] = static_cast(bits >> 16); - } -} - -void GatherPerLayerEmbedding(std::span per_layer_token_embd, - std::int64_t layers, std::int64_t ple_dim, - std::int64_t vocab, std::int64_t token, - std::span out) { - if (layers <= 0 || ple_dim <= 0 || vocab <= 0) { - throw std::runtime_error( - "GatherPerLayerEmbedding layers, ple_dim and vocab must be positive"); - } - if (token < 0 || token >= vocab) { - throw std::runtime_error("token " + std::to_string(token) + - " is outside vocabulary of size " + std::to_string(vocab)); - } - const std::int64_t rows = layers * ple_dim; - if (out.size() != static_cast(rows)) { - throw std::runtime_error("GatherPerLayerEmbedding out width " + - std::to_string(out.size()) + - " does not match layers * ple_dim " + - std::to_string(rows)); - } - - // THE MEMORY LAYOUT, AND IT IS THE OPPOSITE OF WHAT THIS USED TO ASSUME. - // - // `per_layer_token_embd.weight`'s GGUF dims are - // `[layers * ple_dim, vocab]`, and dims[0] is the FASTEST-VARYING extent - // -- so the bytes are `vocab` rows of `layers * ple_dim` contiguous - // elements, with the Q8_0 blocks running along the PLE axis. One token's - // whole per-layer slice is ONE CONTIGUOUS ROW: 280 blocks on E2B, 336 on - // E4B. That is exactly what the reference driver does - // (`gemma4_driver.py`'s `per_layer_embedding()` builds - // `LazyEmbedding(data, dtype, (dims[1], dims[0]))` and slices row i at - // `i * (nbytes // vocab)`), and what `_ple_inputs` then reshapes to - // `[rows, layers, ple_dim]`. - // - // This function previously read the transpose -- `layers * ple_dim` rows - // of `vocab`, one block plucked from each -- because - // verified-gguf-facts.md's "(layers * ple_dim, vocab)" was taken as a - // row-major shape rather than as GGUF dims, and that document said in so - // many words that no change was needed here. THE TWO LAYOUTS HAVE THE - // IDENTICAL TOTAL BYTE COUNT (8960 x 262144 elements either way), so - // every length check passed and the result was 8960 finite, plausible, - // wrong floats. test_gemma4_real_gguf.cpp now pins it against values - // dumped from the reference driver's own rule over the real files. - constexpr std::int64_t kBlockElements = 32; - constexpr std::size_t kBlockBytes = 34; - if (rows % kBlockElements != 0) { - throw std::runtime_error( - "layers * ple_dim is " + std::to_string(rows) + - ", which is not a multiple of 32; the Q8_0 blocks run along THIS " - "axis, not along the vocabulary"); - } - const std::size_t row_bytes = - static_cast(rows / kBlockElements) * kBlockBytes; - const std::size_t required_bytes = static_cast(vocab) * row_bytes; - if (per_layer_token_embd.size() < required_bytes) { - throw std::runtime_error("per_layer_token_embd byte length " + - std::to_string(per_layer_token_embd.size()) + - " is smaller than the " + std::to_string(required_bytes) + - " required for " + std::to_string(vocab) + - " rows of " + std::to_string(rows)); - } - - const auto row = per_layer_token_embd.subspan( - static_cast(token) * row_bytes, row_bytes); - DecodeQ8Row(row, out); -} - -std::vector> BuildPerLayerInputs( - std::span per_layer_token_embd, - std::span model_proj, - std::span proj_norm_gamma, - std::span scaled_embedding, std::span ids, - std::int64_t layers, std::int64_t ple_dim, std::int64_t hidden, - std::int64_t vocab, float epsilon) { - if (layers <= 0 || ple_dim <= 0 || hidden <= 0 || vocab <= 0) { - throw std::runtime_error( - "BuildPerLayerInputs layers, ple_dim, hidden and vocab must be positive"); - } - const auto rows = ids.size(); - if (rows == 0) throw std::runtime_error("BuildPerLayerInputs was given no token ids"); - const auto width = static_cast(layers * ple_dim); - if (proj_norm_gamma.size() != static_cast(ple_dim)) { - throw std::runtime_error("per_layer_proj_norm is " + - std::to_string(proj_norm_gamma.size()) + - " wide, not ple_dim " + std::to_string(ple_dim)); - } - if (model_proj.size() != width * static_cast(hidden)) { - throw std::runtime_error("per_layer_model_proj holds " + - std::to_string(model_proj.size()) + - " elements, not layers * ple_dim * hidden " + - std::to_string(width * static_cast(hidden))); - } - if (scaled_embedding.size() != rows * static_cast(hidden)) { - throw std::runtime_error("the scaled embedding holds " + - std::to_string(scaled_embedding.size()) + - " elements, not rows * hidden " + - std::to_string(rows * static_cast(hidden))); - } - - const auto ple_scale = static_cast(std::sqrt(static_cast(ple_dim))); - const auto projection_scale = - static_cast(1.0 / std::sqrt(static_cast(hidden))); - const auto sum_scale = static_cast(1.0 / std::sqrt(2.0)); - - std::vector> planes( - static_cast(layers), - std::vector(rows * static_cast(ple_dim), 0.0f)); - - // Projected a chunk of rows at a time, so a 4096-token prompt does not hold - // a second [rows, width] FP32 copy beside `planes`. - constexpr std::size_t kRowChunk = 256; - const auto hidden_size = static_cast(hidden); - std::vector gathered(width); - std::vector projected(std::min(rows, kRowChunk) * width); - std::vector normed(static_cast(ple_dim)); - for (std::size_t chunk = 0; chunk < rows; chunk += kRowChunk) { - const std::size_t chunk_rows = std::min(kRowChunk, rows - chunk); - // THE PROJECTION, AGAINST THE TRANSPOSE OF THE MAPPING. Output column - // `c` is the dot product of a row's embedding with mapping ROW `c` - // -- `per_layer_model_proj.weight`'s file dims are [hidden, width] with - // the input width fastest-varying, so the memory is `width` rows of - // `hidden`, and the driver's `embed_rows(ids) @ floats(...)` is against - // `floats()`'s transposed view of it. Reading it the other way round - // has the identical element count. - Project(scaled_embedding.data() + chunk * hidden_size, model_proj.data(), - projected.data(), chunk_rows, width, hidden_size); - - for (std::size_t local = 0; local < chunk_rows; ++local) { - const std::size_t row = chunk + local; - const auto token = static_cast(ids[row]); - // ggml_get_rows, then ggml_scale(sqrt(n_embd_per_layer)). The gather - // reads ONE CONTIGUOUS ROW of the Q8_0 table; see - // GatherPerLayerEmbedding for why the transpose has the identical - // byte count and returns plausible nonsense. - GatherPerLayerEmbedding(per_layer_token_embd, layers, ple_dim, vocab, token, - gathered); - float* projected_row = projected.data() + local * width; - for (std::size_t column = 0; column < width; ++column) - projected_row[column] *= projection_scale; - - for (std::int64_t layer = 0; layer < layers; ++layer) { - const auto offset = static_cast(layer * ple_dim); - // PER LAYER SLICE, not over the whole width: each layer's - // `ple_dim` values are normalized on their own, with the one - // `per_layer_proj_norm` gamma. - RmsNorm(std::span(projected_row + offset, - static_cast(ple_dim)), - proj_norm_gamma, epsilon, normed); - auto& plane = planes[static_cast(layer)]; - for (std::int64_t i = 0; i < ple_dim; ++i) { - const auto index = static_cast(i); - plane[row * static_cast(ple_dim) + index] = - (normed[index] + gathered[offset + index] * ple_scale) * sum_scale; - } - } - } - } - return planes; -} - -void SoftcapLogits(std::span logits, float cap) { - if (!(cap > 0.0f)) { - throw std::runtime_error("logit softcap must be positive, got " + - std::to_string(cap)); - } - for (float& logit : logits) { - logit = std::tanh(logit / cap) * cap; - } -} - -} // namespace flm::gemma4 diff --git a/src/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp b/src/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp deleted file mode 100644 index 9b7c99f5a..000000000 --- a/src/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp +++ /dev/null @@ -1,374 +0,0 @@ -#include "models/gemma4/rai/aie_next/gemma4_rai_shape_plan.hpp" - -#include "models/gemma4/rai/aie_next/gemma4_rai_constants.hpp" -#include "rai/kernel_grid.hpp" - -#include -#include -#include -#include - -namespace flm::gemma4 { -namespace { - -/// \brief the row buckets `ple_bf16` ships a kernel for, taken verbatim from -/// its own doc comment in the pinned corelib.h: "1 for decode, then -/// 64, 128, 256, 512, 1024, 2048, 3072, 4096". This plan also walks -/// these same buckets to query matmul/ssmlp/flat_mha, which is safe -/// even though those ops' own real grids may be coarser subsets -- -/// every corelib pad helper here rounds UP, so an extra query point -/// never produces a wrong answer, only a redundant one. -constexpr std::array kExecutionRows{ - 1, 64, 128, 256, 512, 1024, 2048, 3072, 4096}; - -/// \brief round `rows` up to the row count `ple_bf16` ships a kernel for -/// -/// \note There is NO `ryzenai_corelib_ple_bf16_pad_rows` (or any other ple -/// padding entry point) in the pinned 0.5.0 corelib.h -- confirmed by -/// reading the header, not assumed from the task brief that expected -/// one. `ryzenai_corelib_ple_bf16`'s own doc comment says M "comes -/// from `x`'s own extent and SELECTS THE KERNEL rather than bounding -/// the work" and "a tensor that sizes itself grows to it -/// automatically" -- there is nothing to ask corelib for here, only a -/// fixed table to reproduce. `FLM_CORELIB_FUNCTIONS` -/// (src/include/rai/corelib_api.hpp) accordingly has no ple entry at -/// all, so there is no ABI slot to call even if one wanted to. -std::int64_t RoundToPleBucket(std::int64_t rows) { - for (const auto bucket : kExecutionRows) { - if (rows <= bucket) return bucket; - } - return kExecutionRows.back(); -} - -std::int64_t MatmulRows(const corelib::ShapeGrid& grid, std::int64_t rows, - std::int64_t logical_k, std::int64_t logical_n, - std::uint32_t group, const char* logical_name) { - return corelib::CoveringRows(grid, rows, logical_k, logical_n, group, logical_name); -} - -/// \brief the flat_mha_bf16_desc for one (geometry, cache role) pair -/// -/// \param swa true for the sliding-window geometry, false for full attention -/// \param shared true for a layer that READS another layer's KV cache -/// \note Called four times, once per pair, and the result is never written -/// again. See gemma4_rai_shape_plan.hpp's class comment for why this is -/// a 2x2 rather than two templates the layer loop patches. -ryzenai_corelib_flat_mha_bf16_desc BuildAttentionDesc(const Gemma4Config& config, - bool swa, bool shared) { - ryzenai_corelib_flat_mha_bf16_desc desc{}; - desc.num_heads = config.q_heads; // Gemma4Config::q_heads - desc.kv_num_heads = config.kv_heads; // Gemma4Config::kv_heads - // Sliding layers run at head_dim (256); full-attention layers run at - // global_head_dim (512) -- THIS is the field that makes the two - // geometries resolve DIFFERENT ELFs (see this file's own top comment and - // gemma4_rai_shape_plan.hpp's class doc). - desc.head_size = swa ? config.head_dim : config.global_head_dim; - desc.max_seq = kMaxSequenceLength; // gemma4_rai_constants.hpp, fixed context length - // FULL ROTARY ON BOTH GEOMETRIES. Settled, not open: the reference - // driver sets `rope_dim=head_size` for all four descriptors - // (gemma4_driver.py:1537-1543), and that is what makes the shipped - // `..._256_4096_lc512_scale1` keys -- which carry NO emb_dim segment -- - // resolve. corelib passes no emb_dim when rope_dim == head_size. - // - // `config.json`'s `partial_rotary_factor: 0.25` is NOT a shorter rotary, - // and reading it as one is the trap this comment replaces. It is a - // per-pair FREQUENCY DIVISOR carried in the `rope_freqs.weight` tensor -- - // which is why that tensor is [256] on both E2B and E4B: one entry per - // rotary PAIR of a 512-wide head, not a rotary width of 256. llama.cpp - // passes it as `freq_factors` on FULL-ATTENTION LAYERS ONLY, where it - // divides the frequency (`theta /= freq_factors[i/2]`). - // - // That makes it a fact about the ROPE TABLES (the cos/sin caches, Task - // C7's `_rope_tables` equivalent), which are a separate INPUT to - // flat_mha, not about this field. Two table pairs per model -- base 1e4 - // over the sliding head width with no factors, base 1e6 over the full - // head width with factors -- held for the life of the engine and never - // rewritten in place, because corelib memoizes the derived table on - // buffer identity. Nothing about them changes `rope_dim`. - desc.rope_dim = desc.head_size; - // Sliding layers window at sliding_window (512, from the config); full- - // attention layers see the whole sequence -- 0 is corelib's own "no - // window" (ryzenai_corelib_flat_mha_bf16_desc's doc comment). The window - // is passed UNCONDITIONALLY on every sliding layer: DD decides whether it - // is actually active (`include_lc = 0 < window < sequence`) and a second - // rule here could only disagree with it (gemma4_driver.py:1545-1548). - desc.window = swa ? config.sliding_window : 0; - // kv_shared selects the "_kvshare" ELFs for a layer that reads ANOTHER - // layer's cache (Gemma4Config::layer_cache_owner). It is NOT a cosmetic - // flag: a kvshare kernel "is not handed K at all: it never ropes a K row - // and never writes one" (corelib.h:1772-1779), so getting it wrong in - // either direction runs a different computation over the same buffers and - // returns plausible numbers. The driver's statement of the stakes: "An - // owning layer given a kvshare descriptor never writes its own cache; a - // sharing layer given an owning one overwrites somebody else's. Neither - // errors." (gemma4_driver.py:1526-1559.) - // - // So this is a descriptor axis, materialized here, and NOT a field a - // layer loop patches on a shared template -- which is what this plan used - // to do. See gemma4_rai_shape_plan.hpp's class comment. - desc.kv_shared = shared ? 1 : 0; - // GEMMA 4'S ACTUAL SOFTMAX SCALE, WHICH REALLY IS 1.0. llama.cpp - // hard-codes `f_attention_scale = 1.0f` for this architecture - // (src/models/gemma4.cpp:11) and it reaches the softmax unmodified; the - // reference driver states the same verbatim - // (ryzenai-corelib/python/gemma4_driver.py:239-247). This is not a knob - // and not a value to look up in the config -- it is the model's number, - // and it happens to be the one corelib keys its `_scale1` artifacts on. - // - // THE ENGINE MUST NOT PRE-SCALE Q. corelib's field doc says the `_scale1` - // kernels apply no 1/sqrt(head_size) because "the caller has already - // scaled Q"; for Gemma 4 that caller is THE MODEL'S OWN WEIGHTS. The - // 1/sqrt(head) is carried by the learned QK-Norm gammas (constant-valued - // vectors whose product is 0.125 on a sliding layer and 0.0625 on a full - // one, and it is the K gamma that carries it, not the Q one). THERE IS - // NOTHING FOR THE HOST TO FOLD. A Q pre-scale added on the strength of - // that header sentence would scale attention a second time -- 16x on - // sliding layers, 32x on full ones -- and no test in this suite could see - // it. Reading the header's general statement about the kernel family as a - // statement about this model is precisely the error this comment replaces. - // - // 0.0f, which this field used to hold, means "unset" and selects the - // default (unscaled) family. That family DOES NOT SHIP for Gemma 4 E2B's - // 8/1/256 signature -- corelib.h:1750-1756 names this exact signature as - // shipping ONLY `_scale1` artifacts -- so Build()'s own - // flat_mha_bf16_pad_rows lookup would be rejected on hardware with "no - // attention kernel ships as ...". Same literal on both geometries, which - // is also what the driver does (gemma4_driver.py:1550-1558). - desc.scale = 1.0f; - return desc; -} - -ryzenai_corelib_paged_mha_bf16_desc AsPaged(const ryzenai_corelib_flat_mha_bf16_desc& flat, - std::int64_t context) { - ryzenai_corelib_paged_mha_bf16_desc desc{}; - desc.num_heads = flat.num_heads; - desc.kv_num_heads = flat.kv_num_heads; - desc.head_size = flat.head_size; - desc.max_seq = context; - desc.page_size = kPageSize; - desc.rope_dim = flat.rope_dim; - desc.window = flat.window; - desc.kv_shared = flat.kv_shared; - desc.scale = flat.scale; - return desc; -} - -} // namespace - -Gemma4ShapePlan Gemma4ShapePlan::Build( - const std::shared_ptr& api, - ryzenai_corelib_stream_ptr stream, const Gemma4Config& config, - std::int64_t context) { - if (!api) throw std::invalid_argument("Gemma4ShapePlan corelib API is null"); - if (!stream) throw std::invalid_argument("Gemma4ShapePlan stream is null"); - - Gemma4ShapePlan plan; - // All four (geometry, cache role) descriptors, up front. Nothing below - // writes to any of them again. - for (const bool swa : {true, false}) { - for (const bool shared : {false, true}) { - plan.attention_descs_[DescIndex(swa, shared)] = - BuildAttentionDesc(config, swa, shared); - } - } - plan.paged_ = context > kPageSize; - if (plan.paged_) { - if (context > kMaxPagedContext || context % kPageSize != 0) { - throw std::invalid_argument( - "Gemma 4 paged context must be a multiple of 4096 up to " + - std::to_string(kMaxPagedContext)); - } - for (std::size_t index = 0; index < plan.attention_descs_.size(); ++index) { - plan.paged_descs_[index] = AsPaged(plan.attention_descs_[index], context); - } - } - - // ssmlp must be planned at EVERY DISTINCT width in - // config.layer_intermediate, not just the first layer's -- E2B steps - // from 6144 to 12288 at layer 15, so a plan built from layer 0 alone - // would be wrong for layers 15-34. E4B has exactly one width; the loop - // below still runs once for it. - std::vector widths; - for (const auto width : config.layer_intermediate) { - if (std::find(widths.begin(), widths.end(), width) == widths.end()) { - widths.push_back(width); - } - } - if (widths.empty()) { - throw std::invalid_argument("Gemma4ShapePlan: config has no FFN widths"); - } - - plan.sliding_rows_.reserve(static_cast(kMaxSequenceLength)); - plan.full_rows_.reserve(static_cast(kMaxSequenceLength)); - const auto matmul = corelib::ShapeGrid::Matmul(*api, stream); - const auto ssmlp_grid = corelib::ShapeGrid::SsMlp(*api, stream, true); - // A multi-row pass at position 0 is attention over the first page, and - // corelib's paged_mha runs it on flat_mha's own prefill kernel at the - // pass's rows; every later chunk is a whole page on the paged kernel. So - // the row extents come from flat_mha's grid even on a paged cache, and - // the paged grid only has to carry the one page-sized chunk. - const auto sliding_mha = - corelib::ShapeGrid::FlatMha(*api, stream, plan.attention_descs_[DescIndex(true, false)]); - const auto full_mha = - corelib::ShapeGrid::FlatMha(*api, stream, plan.attention_descs_[DescIndex(false, false)]); - if (plan.paged_) { - for (const bool swa : {true, false}) { - const auto paged = corelib::ShapeGrid::PagedMha( - *api, stream, plan.paged_descs_[DescIndex(swa, false)]); - const char* what = swa ? "sliding paged_mha chunk" : "full paged_mha chunk"; - if (corelib::CoveringRows(paged, kPageSize, 0, 0, -1, what) != kPageSize) { - throw std::runtime_error(std::string(what) + " has no kernel at exactly " + - std::to_string(kPageSize) + " rows"); - } - } - } - - for (const auto rows : kExecutionRows) { - // ssmlp: gelu (activation = 1) with the post-feedforward norm SET - // (post_feedforward_layernorm = 1). - // - // `activation` is a DECLARATION, not a switch: it is baked into the - // ELF, appears nowhere in the artifact's name, and on the no-bias - // path the runtime parameter DD derives from it is never written to - // the weights BO -- so setting it wrongly would not change what the - // NPU computes, only make the descriptor lie about what it computes. - // Set correctly anyway: it is the only place the block's activation - // is written down. NEVER "debug" by flipping it. - // - // `post_feedforward_layernorm` genuinely changes what the block - // computes and writes: set, the norm moves inside the residual and - // the `normalized` output plane is never written (though corelib - // still requires a distinct buffer for that port) -- this is the - // field that keys the "gemma_fusion*" ELFs instead of - // "fusion_mladf_ssmlp*". See verified-gguf-facts.md's "CONSTRAINT - // FOR GEMMA 4'S ENGINE" for the full data-flow consequence (Task - // C8's problem, not this one's -- this plan only asks how such a - // descriptor pads, it never dispatches ssmlp). - std::int64_t ssmlp_rows = rows; - for (const auto width : widths) { - ssmlp_rows = std::max( - ssmlp_rows, corelib::CoveringRows(ssmlp_grid, rows, config.hidden, width, - config.group, "ssmlp")); - } - - const auto ple_rows = RoundToPleBucket(rows); - // `ryzenai_corelib_rmsnorm_bf16_pad_rows` exists in corelib.h, but is - // documented as an IDENTITY FUNCTION that does NOT report the - // covering extent -- "sizing a buffer from its return value is the - // trap this paragraph exists to name" (corelib.h's own - // ryzenai_corelib_rmsnorm_bf16 doc). It is also not part of - // FLM_CORELIB_FUNCTIONS (src/include/rai/corelib_api.hpp), so there - // is no binding to call it through even if its answer were - // trustworthy. Every rmsnorm in a Gemma 4 layer runs over the same - // per-token residual stream `ple_bf16` already sizes at this bucket, - // so its already-safe covering extent is reused here rather than - // inventing a second, uncheckable source of truth. - const auto rmsnorm_rows = ple_rows; - - for (const bool swa : {true, false}) { - // The OWNING descriptor of this geometry, and that choice is not - // arbitrary: the row extents this loop fills are per GEOMETRY - // only, never per (geometry, cache role). The header says so as - // the stated reason flat_mha_bf16_elf_name exists at all -- the - // descriptor's "window, kv_shared and scale do not change the - // shapes of anything -- they change WHICH KERNEL RUNS, by - // changing the name DD resolves" (corelib.h:1810-1812). So - // flat_mha_bf16_pad_rows cannot answer differently for the - // kvshare family, and sliding_rows_/full_rows_ are two tables, - // not four. Do not "fix" this into a four-way split. - const auto& desc = plan.attention_descs_[DescIndex(swa, /*shared=*/false)]; - // EVERY PROJECTION -- Q, K, V AND the output -- is packed at - // `config.group`. Not a judgement call: the reference driver - // declares the field as "group: int # every projection" - // (gemma4_driver.py:537-539), its `_matmul_weights` defaults - // `group_size=self.cfg.group` (:1383-1387), and the ONLY call in - // the whole driver that overrides that default is lm_head - // (:1691-1694). Its enumeration of a prefill's dispatches puts - // all three shapes at g32 == `cfg.group`: "Q/K/O at 1536->2048, - // 1536->256 and 2048->1536 on a sliding layer and 1536->4096, - // 1536->512 and 4096->1536 on a full one" (:294-298). - // - // `config.head_group` is lm_head's, and lm_head's alone -- see - // the single query after this loop. It is a SEPARATE FIELD - // because on other models the two genuinely differ; on both - // shipped Gemma 4 rows they are both 32 (gemma4_rai_gguf.cpp), - // which is exactly why a version of this code that queried Q and - // K/V at `head_group` passed every test in the suite. Do not - // reintroduce it here on the strength of "head"-shaped outputs: - // `head_group` is named for lm_head, not for the attention heads. - Gemma4RowExtents extents{}; - extents.q_rows = MatmulRows(matmul, rows, config.hidden, - desc.num_heads * desc.head_size, - config.group, "query"); - extents.kv_rows = MatmulRows(matmul, rows, config.hidden, - desc.kv_num_heads * desc.head_size, - config.group, "key/value"); - extents.output_rows = MatmulRows(matmul, rows, - desc.num_heads * desc.head_size, - config.hidden, config.group, - "output"); - extents.ssmlp_rows = ssmlp_rows; - extents.rmsnorm_rows = rmsnorm_rows; - extents.ple_rows = ple_rows; - - const auto& mha_grid = swa ? sliding_mha : full_mha; - extents.flat_mha_rows = corelib::CoveringRows(mha_grid, rows, 0, 0, -1, "flat_mha"); - - auto& bucket_rows = swa ? plan.sliding_rows_ : plan.full_rows_; - while (bucket_rows.size() < static_cast(rows)) { - bucket_rows.push_back(extents); - } - } - } - - // LM_HEAD, AND IT IS THE ONE MATMUL OUTSIDE THE BUCKET LOOP ABOVE. - // - // Two things make it different from every other matmul in the model, and - // both come from the reference driver rather than from this port: - // - // - IT HAS EXACTLY ONE ROW SHAPE. "ONE ROW IS NOT A SIMPLIFICATION, IT - // IS THE ONLY SHAPE AVAILABLE. lm_head is 1536 -> 262144 and ships at - // M == 1 alone at this model's (8, 17) pair, so there is no batched - // form to offer even if a caller wanted every position's logits" - // (gemma4_driver.py:2449-2453). The driver says the same from the - // prefill side -- "lm_head runs at M=1, so neither is a prefill - // matmul" (:294-298). Walking it through `kExecutionRows` would - // therefore ask corelib about shapes that do not exist; at best the - // answers are meaningless, at worst the query itself fails on - // hardware. - // - IT IS THE ONE WEIGHT PACKED AT `head_group`. "TIED to token_embd -- - // this conversion has no output.weight -- and at `head_group`, which - // on this model happens to equal `group`" (:1691-1694); "the only - // weight packed at `cfg.head_group`" (:2449-2453). - // - // It is asked of corelib rather than hardcoded to 1 because the answer is - // corelib's to give, and because a fake or a future corelib that answers - // otherwise should show up here rather than in C7's allocation. - plan.lm_head_rows_ = MatmulRows(matmul, /*rows=*/1, config.hidden, - config.vocab, config.head_group, "lm_head"); - - return plan; -} - -const Gemma4RowExtents& Gemma4ShapePlan::ForRows(std::size_t live_rows, bool swa) const { - const auto& bucket_rows = swa ? sliding_rows_ : full_rows_; - if (live_rows == 0 || live_rows > bucket_rows.size()) { - throw std::out_of_range("Gemma 4 live rows must be in 1.." + - std::to_string(bucket_rows.size())); - } - return bucket_rows[live_rows - 1]; -} - -const ryzenai_corelib_flat_mha_bf16_desc& -Gemma4ShapePlan::attention_desc(bool swa, bool shared) const noexcept { - return attention_descs_[DescIndex(swa, shared)]; -} - -const ryzenai_corelib_paged_mha_bf16_desc& -Gemma4ShapePlan::paged_attention_desc(bool swa, bool shared) const noexcept { - return paged_descs_[DescIndex(swa, shared)]; -} - -std::int64_t Gemma4ShapePlan::lm_head_rows() const noexcept { return lm_head_rows_; } - -} // namespace flm::gemma4 diff --git a/src/common/models/qwen3/rai/aie_next/qwen3_rai.cpp b/src/common/models/qwen3/rai/aie_next/qwen3_rai.cpp deleted file mode 100644 index 23810a576..000000000 --- a/src/common/models/qwen3/rai/aie_next/qwen3_rai.cpp +++ /dev/null @@ -1,667 +0,0 @@ -#include "models/qwen3/rai/aie_next/qwen3_rai.hpp" - -#include "models/qwen3/rai/aie_next/qwen3_rai_shape_plan.hpp" -#include "rai/corelib_object.hpp" -#include "rai/host_ops.hpp" -#include "rai/weight_cache.hpp" -#include "rai/weight_source.hpp" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace flm::qwen3 { -namespace { -using namespace flm::corelib; - -std::string Name(std::size_t i, const char* suffix) { return "blk." + std::to_string(i) + suffix; } - -/// \brief every packing parameter the cached bytes depend on -std::string CacheLayout(const Qwen3Config& c, bool tied) { - std::ostringstream out; - out << "qwen3-" << c.size << " L" << c.layers << " h" << c.hidden << " q" << c.q_dim() - << " kv" << c.kv_dim() << " i" << c.intermediate << " v" << c.vocab << " g" << c.group - << " hg" << c.head_group << (tied ? " tied" : " untied"); - return out.str(); -} - -/// \brief load-time phase accounting, printed when FLM_RAI_PROFILE_LOAD is set -struct LoadPhases { - std::chrono::steady_clock::time_point mark{std::chrono::steady_clock::now()}; - double host_prep{}, shape_plan{}, weight_create{}, device_tensors{}, cache_write{}; - std::uint64_t cache_reclaimed{}; - bool from_cache{}; - - double Lap() { - const auto now = std::chrono::steady_clock::now(); - const double seconds = std::chrono::duration(now - mark).count(); - mark = now; - return seconds; - } - - void Report() const { - const char* enabled = std::getenv("FLM_RAI_PROFILE_LOAD"); - if (!enabled || !*enabled || *enabled == '0') return; - std::ostringstream out; - out << std::fixed << std::setprecision(2) << "[FLM] rai load: " - << host_prep + shape_plan + weight_create + device_tensors + cache_write - << " s total (host prep " << host_prep << ", shape plan " << shape_plan - << ", weight " << (from_cache ? "cache" : "requantize") << " " << weight_create - << ", device tensors " << device_tensors << ")"; - if (cache_write > 0.0) out << " [cache written in " << cache_write << " s]"; - if (cache_reclaimed > 0) - out << " [reclaimed " << cache_reclaimed / (1024 * 1024) << " MB of stale cache]"; - std::cout << out.str() << std::endl; - } -}; -} // namespace - -struct qwen3_rai::Impl { - std::shared_ptr package; - std::shared_ptr runtime; - std::shared_ptr api; - const Qwen3Config* c{}; - LoadPhases phases; - Qwen3ShapePlan plan; - std::uint32_t max_length; - int position{}; - std::optional saved; - /// \brief the tokens whose K/V rows fill [0, history.size()) - std::vector history; - bool poisoned{}; - UniqueStream stream; - std::vector q_weights, k_weights, v_weights, o_weights; - std::vector mlp_weights; - std::vector q_norm_weights, k_norm_weights; - UniqueMatMulWeights lm_weights; - // flat_mha reads the rotary tables on the host and corelib copies nothing, - // so this storage must outlive the views below -- hence declared first. - std::vector rope_cosine, rope_sine; - UniqueHostView cosine, sine; - UniqueTensor hidden, residual, skip, q, k, q_normed, k_normed, attention, lm_input, logits; - std::vector k_cache, v_cache; - - Impl(std::shared_ptr pkg, std::shared_ptr rt, - std::uint32_t maximum) - : package(std::move(pkg)), runtime(std::move(rt)), - api(runtime ? runtime->api() : nullptr), max_length(maximum) { - if (!package) throw std::invalid_argument("Qwen3 GGUF package is null"); - if (!runtime || !api) throw std::invalid_argument("corelib runtime is null"); - if (!maximum || maximum > kMaxSequenceLength) - throw std::invalid_argument("Qwen3 maximum length must be in 1..4096"); - c = &package->Config(); - const auto layers = static_cast(c->layers); - - std::vector> attn_norm(layers), ffn_norm(layers), q_norm(layers), - k_norm(layers); - for (std::size_t i = 0; i < layers; ++i) { - const auto& t = package->Layer(i); - attn_norm[i] = rai::F32ToBf16(t.attn_norm.values); - ffn_norm[i] = rai::F32ToBf16(t.ffn_norm.values); - q_norm[i] = rai::F32ToBf16(t.q_norm.values); - k_norm[i] = rai::F32ToBf16(t.k_norm.values); - } - const auto final_norm = rai::F32ToBf16(package->OutputNorm().values); - const std::array epsilon_f32{c->epsilon}; - const auto epsilon = rai::F32ToBf16(epsilon_f32); - BuildRopeTables(); - phases.host_prep = phases.Lap(); - - auto lease = runtime->AcquireExecution(); - stream = corelib::OpenModelStream(api, kPrefillPdi, kTokenPdi, "qwen3"); - plan = Qwen3ShapePlan::Build(api, stream.get(), *c); - phases.shape_plan = phases.Lap(); - - q_weights.resize(layers); k_weights.resize(layers); v_weights.resize(layers); - o_weights.resize(layers); mlp_weights.resize(layers); - q_norm_weights.resize(layers); k_norm_weights.resize(layers); - - // QK-Norm scales go through the reference packer: they are 128 wide, - // cost nothing to pack, and so are not worth a place in the cache. - const ryzenai_corelib_rmsnorm_bf16_weights_desc norm_desc{kQkNormWidth, c->epsilon}; - const auto norm = [&](const std::vector& scale, const std::string& label) { - const ryzenai_corelib_weights_source sources[]{ - corelib::Bf16(ryzenai_corelib_weights_role_scale, scale.data(), scale.size())}; - void* p = nullptr; - api->Check(api->functions().rmsnorm_weights_create(&norm_desc, sources, std::size(sources), - nullptr, &p), - "ryzenai_corelib_rmsnorm_bf16_weights_create " + label); - return UniqueRmsNormWeights(api, p); - }; - for (std::size_t i = 0; i < layers; ++i) { - q_norm_weights[i] = norm(q_norm[i], Name(i, ".attn_q_norm")); - k_norm_weights[i] = norm(k_norm[i], Name(i, ".attn_k_norm")); - } - - const auto matmul = [&](const rai::GgufTensor& tensor, std::int64_t k_in, std::int64_t n_out, - std::uint32_t group, const std::string& label) { - const ryzenai_corelib_matmul_bf16_weights_desc desc{k_in, n_out, group, false}; - const ryzenai_corelib_weights_source sources[]{corelib::GgufQ8( - ryzenai_corelib_weights_role_qweight, tensor.bytes.data(), tensor.bytes.size(), - n_out, k_in)}; - const auto options = corelib::PackOptions(kRequantizeThreads); - void* p = nullptr; - api->Check(api->functions().matmul_weights_create( - &desc, sources, std::size(sources), &options, &p), - "ryzenai_corelib_matmul_bf16_weights_create " + label); - return UniqueMatMulWeights(api, p); - }; - std::vector> creates; - creates.reserve(layers * 5 + 1); - for (std::size_t i = 0; i < layers; ++i) { - creates.push_back([&, i] { q_weights[i] = matmul(package->Layer(i).q, c->hidden, c->q_dim(), c->group, Name(i, ".attn_q")); }); - creates.push_back([&, i] { k_weights[i] = matmul(package->Layer(i).k, c->hidden, c->kv_dim(), c->group, Name(i, ".attn_k")); }); - creates.push_back([&, i] { v_weights[i] = matmul(package->Layer(i).v, c->hidden, c->kv_dim(), c->group, Name(i, ".attn_v")); }); - creates.push_back([&, i] { o_weights[i] = matmul(package->Layer(i).o, c->q_dim(), c->hidden, c->group, Name(i, ".attn_output")); }); - creates.push_back([&, i] { - const auto& t = package->Layer(i); - // The fused MLP applies the NEXT block's input norm to its own - // output, so norm1 is layer i+1's attn_norm, or output_norm - // after the last layer -- which is why nothing applies a final - // norm before lm_head. - const auto& next = i + 1 < layers ? attn_norm[i + 1] : final_norm; - const ryzenai_corelib_ssmlp_bf16_weights_desc desc{c->hidden, c->intermediate, c->group}; - const ryzenai_corelib_weights_source sources[]{ - corelib::Bf16(ryzenai_corelib_weights_role_epsilon, epsilon.data(), epsilon.size()), - corelib::Bf16(ryzenai_corelib_weights_role_norm0, ffn_norm[i].data(), ffn_norm[i].size()), - corelib::Bf16(ryzenai_corelib_weights_role_norm1, next.data(), next.size()), - corelib::GgufQ8(ryzenai_corelib_weights_role_gate_qweight, t.gate.bytes.data(), t.gate.bytes.size(), c->intermediate, c->hidden), - corelib::GgufQ8(ryzenai_corelib_weights_role_up_qweight, t.up.bytes.data(), t.up.bytes.size(), c->intermediate, c->hidden), - corelib::GgufQ8(ryzenai_corelib_weights_role_down_qweight, t.down.bytes.data(), t.down.bytes.size(), c->hidden, c->intermediate)}; - const auto options = corelib::PackOptions(kRequantizeThreads); - void* p = nullptr; - api->Check(api->functions().ssmlp_weights_create( - &desc, sources, std::size(sources), &options, &p), - "ryzenai_corelib_ssmlp_bf16_weights_create layer " + - std::to_string(i)); - mlp_weights[i] = UniqueSsMlpWeights(api, p); - }); - } - creates.push_back([&] { - lm_weights = matmul(package->LmHead(), c->hidden, c->vocab, c->head_group, - std::string(package->LmHead().name)); - }); - - const auto cache = rai::WeightCache::ForGguf(package->Path()); - const auto version = api->runtime_version(); - const auto key = rai::MakeWeightCacheKey(package->Path(), version.major, version.minor, - version.patch, - CacheLayout(*c, package->TiedLmHead()), - creates.size()); - bool from_cache = false; - if (cache) { - if (const auto spans = cache->ReadIndex(key)) - from_cache = LoadWeightsFromCache(cache->DataPath(), *spans); - } - if (!from_cache) { - std::cout << "[FLM] Packing " << creates.size() << " Qwen3-" << c->size - << " weights from Q8_0" << (cache ? "; cached for the next load" : "") - << std::endl; - RunCreates(creates); - } - phases.weight_create = phases.Lap(); - phases.from_cache = from_cache; - if (cache && !from_cache) { - // Whatever is there did not match, or it would have been used. - phases.cache_reclaimed = cache->Remove(); - WriteWeightCache(*cache, key); - phases.cache_write = phases.Lap(); - } - - const auto tensor = [&](std::initializer_list dims, const char* label) { - std::vector shape(dims); - void* p = nullptr; - api->Check(api->functions().create_device_tensor(ryzenai_corelib_data_type_bf16, - shape.data(), shape.size(), &p), - std::string("ryzenai_corelib_create_device_tensor ") + label); - return UniqueTensor(api, p); - }; - constexpr auto rows = kMaxSequenceLength; - hidden = tensor({rows, c->hidden}, "hidden"); - residual = tensor({rows, c->hidden}, "residual"); - skip = tensor({rows, c->hidden}, "skip"); - q = tensor({rows, c->q_dim()}, "query"); - k = tensor({rows, c->kv_dim()}, "key"); - // QK-Norm's destinations. Not in place: at k=128 an in-place RMSNorm - // corrupts a short run of rows for M >= 1024, on exact kernels too. - q_normed = tensor({rows, c->q_dim()}, "normed query"); - k_normed = tensor({rows, c->kv_dim()}, "normed key"); - attention = tensor({rows, c->q_dim()}, "attention"); - // lm_head runs one row; a one-row tensor is what selects the token PDI. - lm_input = tensor({1, c->hidden}, "lm input"); - logits = tensor({1, c->vocab}, "logits"); - k_cache.resize(layers); - v_cache.resize(layers); - for (std::size_t i = 0; i < layers; ++i) { - k_cache[i] = tensor({c->kv_heads, rows, c->head_dim}, "K cache"); - v_cache[i] = tensor({c->kv_heads, rows, c->head_dim}, "V cache"); - } - const auto host_view = [&](const std::vector& values, const char* label) { - const std::array shape{kMaxSequenceLength, c->head_dim / 2}; - void* p = nullptr; - api->Check(api->functions().create_host_view(ryzenai_corelib_data_type_fp32, shape.data(), - shape.size(), values.data(), &p), - std::string("ryzenai_corelib_create_host_view ") + label); - return UniqueHostView(api, p); - }; - cosine = host_view(rope_cosine, "cosine"); - sine = host_view(rope_sine, "sine"); - phases.device_tensors = phases.Lap(); - phases.Report(); - } - - /// \brief cos/sin over the whole head, from the file's own rope theta - /// \note Qwen3 rotates all of head_dim and ships no rope factors, so theta - /// is the whole derivation. Double precision throughout; float32 - /// drifts measurably at the far end of the table. - void BuildRopeTables() { - const auto half = static_cast(c->head_dim / 2); - rope_cosine.resize(static_cast(kMaxSequenceLength) * half); - rope_sine.resize(rope_cosine.size()); - for (std::size_t i = 0; i < half; ++i) { - const double inv_freq = - 1.0 / std::pow(c->rope_theta, 2.0 * static_cast(i) / static_cast(c->head_dim)); - for (std::size_t p = 0; p < static_cast(kMaxSequenceLength); ++p) { - const double angle = static_cast(p) * inv_freq; - rope_cosine[p * half + i] = static_cast(std::cos(angle)); - rope_sine[p * half + i] = static_cast(std::sin(angle)); - } - } - } - - void RunCreates(const std::vector>& creates) { - std::atomic next{0}; - std::mutex failure_mutex; - std::exception_ptr first_failure; - const auto work = [&] { - for (std::size_t i = next++; i < creates.size(); i = next++) { - try { - creates[i](); - } catch (...) { - std::lock_guard lock(failure_mutex); - if (!first_failure) first_failure = std::current_exception(); - next = creates.size(); - return; - } - } - }; - const auto workers = std::min(kWeightCreateConcurrency, creates.size()); - std::vector pool; - for (std::size_t t = 1; t < workers; ++t) pool.emplace_back(work); - work(); - for (auto& worker : pool) worker.join(); - if (first_failure) std::rethrow_exception(first_failure); - } - - /// \note Slots are five per layer (Q, K, V, O, MLP) then lm_head. Both cache - /// directions walk this order, so an index entry always refers to the - /// weight it was written from. - std::size_t HeadSlot() const { return static_cast(c->layers) * 5; } - bool SlotIsMatmul(std::size_t slot) const { return slot == HeadSlot() || slot % 5 != 4; } - - ryzenai_corelib_matmul_bf16_weights_desc MatmulDescAt(std::size_t slot) const { - if (slot == HeadSlot()) return {c->hidden, c->vocab, c->head_group, false}; - switch (slot % 5) { - case 0: return {c->hidden, c->q_dim(), c->group, false}; - case 1: - case 2: return {c->hidden, c->kv_dim(), c->group, false}; - default: return {c->q_dim(), c->hidden, c->group, false}; - } - } - - void* WeightAt(std::size_t slot) const { - if (slot == HeadSlot()) return lm_weights.get(); - const auto layer = slot / 5; - switch (slot % 5) { - case 0: return q_weights[layer].get(); - case 1: return k_weights[layer].get(); - case 2: return v_weights[layer].get(); - case 3: return o_weights[layer].get(); - default: return mlp_weights[layer].get(); - } - } - - void AssignWeightAt(std::size_t slot, void* handle) { - if (slot == HeadSlot()) { lm_weights = UniqueMatMulWeights(api, handle); return; } - const auto layer = slot / 5; - switch (slot % 5) { - case 0: q_weights[layer] = UniqueMatMulWeights(api, handle); break; - case 1: k_weights[layer] = UniqueMatMulWeights(api, handle); break; - case 2: v_weights[layer] = UniqueMatMulWeights(api, handle); break; - case 3: o_weights[layer] = UniqueMatMulWeights(api, handle); break; - default: mlp_weights[layer] = UniqueSsMlpWeights(api, handle); break; - } - } - - void ReleaseCachedWeights() { - for (std::size_t slot = 0; slot <= HeadSlot(); ++slot) AssignWeightAt(slot, nullptr); - } - - /// \return true when every slot bound from the file; false leaves none bound - /// \note corelib rejects a slice that is not exactly what its descriptor - /// packs to, so a stale file is caught here, and a half-cached model - /// is never kept: the caller simply packs. - bool LoadWeightsFromCache(const std::filesystem::path& data_path, - const std::vector& spans) { - const auto path = data_path.string(); - if (spans.size() != HeadSlot() + 1) return false; - for (std::size_t slot = 0; slot < spans.size(); ++slot) { - void* handle = nullptr; - ryzenai_corelib_status status; - const auto packed = corelib::PackedFile(path.c_str(), spans[slot].offset, spans[slot].size); - if (SlotIsMatmul(slot)) { - const auto desc = MatmulDescAt(slot); - status = api->functions().matmul_weights_create(&desc, &packed, 1, nullptr, &handle); - } else { - const ryzenai_corelib_ssmlp_bf16_weights_desc desc{c->hidden, c->intermediate, c->group}; - status = api->functions().ssmlp_weights_create(&desc, &packed, 1, nullptr, &handle); - } - if (status != ryzenai_corelib_status_success || handle == nullptr) { - ReleaseCachedWeights(); - return false; - } - AssignWeightAt(slot, handle); - } - return true; - } - - /// \note Best effort: a cache that cannot be written only means the next - /// launch packs again. The index goes last, so a data file without a - /// matching index is never used. - void WriteWeightCache(const rai::WeightCache& cache, const rai::WeightCacheKey& key) { - try { - std::error_code error; - std::filesystem::create_directories(cache.DataPath().parent_path(), error); - const auto data_path = cache.DataPath(); - const auto temporary = data_path.string() + ".tmp"; - std::vector spans; - { - std::ofstream output(temporary, std::ios::binary | std::ios::trunc); - if (!output) return; - std::vector buffer; - std::uint64_t offset = 0; - for (std::size_t slot = 0; slot <= HeadSlot(); ++slot) { - std::size_t size = 0; - if (api->functions().weights_copy_data(WeightAt(slot), nullptr, 0, &size) != - ryzenai_corelib_status_success || size == 0) - return; - buffer.resize(size); - if (api->functions().weights_copy_data(WeightAt(slot), buffer.data(), - buffer.size(), &size) != - ryzenai_corelib_status_success) - return; - output.write(buffer.data(), static_cast(size)); - if (!output) return; - spans.push_back({offset, static_cast(size)}); - offset += size; - } - } - std::filesystem::rename(temporary, data_path, error); - if (error) { std::filesystem::remove(temporary, error); return; } - if (!cache.WriteIndex(key, spans)) std::filesystem::remove(data_path, error); - } catch (...) { - // Caching is an optimisation; never let it fail a load. - } - } - - void usable() const { - if (poisoned) throw std::runtime_error("Qwen3 corelib engine is poisoned"); - } - - /// \brief one pass of `ids` at the current position - buffer Run(std::span ids) { - usable(); - const auto live = ids.size(); - if (live == 0) throw std::invalid_argument("Qwen3 request contains no token IDs"); - if (live > 1 && position != 0) - throw std::logic_error("Qwen3 multi-row pass must start at position zero"); - if (position + live > max_length || position + live > kMaxSequenceLength) - throw std::out_of_range("Qwen3 request exceeds configured context capacity"); - if (live == 1 && position + 1 > kMaxDecodeWindow) - throw std::out_of_range("Qwen3 decode window stops at position 4095"); - - const auto rows = plan.RowsFor(live); - const auto& layer0 = package->Layer(0); - const auto decoded = rai::DecodeQ8Rows(package->Embedding(), ids); - std::vector normalized(decoded.size()); - rai::RmsNorm(decoded, layer0.attn_norm.values, static_cast(live), c->hidden, - c->epsilon, normalized); - std::vector input(static_cast(rows * c->hidden), 0.0f); - std::vector residual_input(input.size(), 0.0f); - std::copy(normalized.begin(), normalized.end(), input.begin()); - std::copy(decoded.begin(), decoded.end(), residual_input.begin()); - - auto lease = runtime->AcquireExecution(); - bool submitted = false; - try { - const auto& f = api->functions(); - api->Check(f.tensor_write(hidden.get(), ryzenai_corelib_data_type_fp32, input.data(), - input.size(), 0), - "ryzenai_corelib_tensor_write hidden"); - api->Check(f.tensor_write(residual.get(), ryzenai_corelib_data_type_fp32, - residual_input.data(), residual_input.size(), 0), - "ryzenai_corelib_tensor_write residual embedding"); - - // Each op reads M off its operand's shape, so every buffer is seen - // through a window at this pass's row count. Built once per pass. - const auto view = [&](void* parent, std::int64_t r, std::int64_t cols, const char* label) { - const std::array shape{r, cols}; - void* p = nullptr; - api->Check(f.create_tensor_window(parent, shape.data(), shape.size(), 0, nullptr, 0, &p), - std::string("ryzenai_corelib_create_tensor_window ") + label); - return UniqueTensorWindow(api, p); - }; - const auto hidden_w = view(hidden.get(), rows, c->hidden, "hidden"); - const auto q_w = view(q.get(), rows, c->q_dim(), "query"); - const auto k_w = view(k.get(), rows, c->kv_dim(), "key"); - const auto q_attn = view(q_normed.get(), rows, c->q_dim(), "normed query"); - const auto k_attn = view(k_normed.get(), rows, c->kv_dim(), "normed key"); - const auto attention_w = view(attention.get(), rows, c->q_dim(), "attention"); - // The per-head view QK-Norm runs over: [rows, heads * 128] read as - // [rows * heads, 128] is the same row-major memory. - const auto q_heads_in = view(q.get(), rows * c->q_heads, c->head_dim, "query heads"); - const auto k_heads_in = view(k.get(), rows * c->kv_heads, c->head_dim, "key heads"); - const auto q_heads_out = view(q_normed.get(), rows * c->q_heads, c->head_dim, "normed query heads"); - const auto k_heads_out = view(k_normed.get(), rows * c->kv_heads, c->head_dim, "normed key heads"); - const auto residual_w = view(residual.get(), rows, c->hidden, "residual"); - const auto skip_w = view(skip.get(), rows, c->hidden, "skip"); - void* res = residual_w.get(); - void* sk = skip_w.get(); - - for (std::size_t i = 0; i < static_cast(c->layers); ++i) { - const auto layer = std::to_string(i); - const auto query_status = f.matmul(stream.get(), hidden_w.get(), q_weights[i].get(), q_w.get()); - submitted = submitted || query_status == ryzenai_corelib_status_success || - query_status == ryzenai_corelib_status_failure; - api->Check(query_status, "ryzenai_corelib_matmul_bf16 query layer " + layer); - api->Check(f.matmul(stream.get(), hidden_w.get(), k_weights[i].get(), k_w.get()), - "ryzenai_corelib_matmul_bf16 key layer " + layer); - // v_proj writes the V cache directly: a window positioned at - // this pass's first row, whose height bounds what lands there. - const std::array v_shape{c->kv_heads, kMaxSequenceLength - position, - c->head_dim}; - const std::array v_strides{kMaxSequenceLength * c->head_dim, - c->head_dim, 1}; - void* p = nullptr; - api->Check(f.create_tensor_window(v_cache[i].get(), v_shape.data(), v_shape.size(), - static_cast(position * c->head_dim), - v_strides.data(), v_strides.size(), &p), - "ryzenai_corelib_create_tensor_window V"); - const UniqueTensorWindow v_window(api, p); - api->Check(f.matmul(stream.get(), hidden_w.get(), v_weights[i].get(), v_window.get()), - "ryzenai_corelib_matmul_bf16 value layer " + layer); - // QK-Norm before attention: flat_mha applies the rotary itself. - api->Check(f.rmsnorm(stream.get(), q_heads_in.get(), q_norm_weights[i].get(), q_heads_out.get()), - "ryzenai_corelib_rmsnorm_bf16 q-norm layer " + layer); - api->Check(f.rmsnorm(stream.get(), k_heads_in.get(), k_norm_weights[i].get(), k_heads_out.get()), - "ryzenai_corelib_rmsnorm_bf16 k-norm layer " + layer); - api->Check(f.flat_mha(stream.get(), &plan.attention_desc(), q_attn.get(), k_attn.get(), - position, cosine.get(), sine.get(), k_cache[i].get(), - v_cache[i].get(), attention_w.get()), - "ryzenai_corelib_flat_mha_bf16 layer " + layer); - api->Check(f.matmul(stream.get(), attention_w.get(), o_weights[i].get(), hidden_w.get()), - "ryzenai_corelib_matmul_bf16 output layer " + layer); - api->Check(f.ssmlp(stream.get(), hidden_w.get(), res, mlp_weights[i].get(), sk, hidden_w.get()), - "ryzenai_corelib_ssmlp_bf16 layer " + layer); - std::swap(res, sk); - } - api->Check(f.stream_synchronize(stream.get()), "ryzenai_corelib_stream_synchronize hidden"); - std::vector row(static_cast(c->hidden)); - api->Check(f.tensor_read(hidden.get(), ryzenai_corelib_data_type_bf16, row.data(), row.size(), - (live - 1) * static_cast(c->hidden)), - "ryzenai_corelib_tensor_read final hidden row"); - api->Check(f.tensor_write(lm_input.get(), ryzenai_corelib_data_type_bf16, row.data(), row.size(), 0), - "ryzenai_corelib_tensor_write LM head input"); - api->Check(f.matmul(stream.get(), lm_input.get(), lm_weights.get(), logits.get()), - "ryzenai_corelib_matmul_bf16 LM head"); - api->Check(f.stream_synchronize(stream.get()), "ryzenai_corelib_stream_synchronize logits"); - buffer out(static_cast(c->vocab)); - api->Check(f.tensor_read(logits.get(), ryzenai_corelib_data_type_bf16, out.data(), out.size(), 0), - "ryzenai_corelib_tensor_read logits"); - position += static_cast(live); - return out; - } catch (...) { - if (submitted) { - (void)api->functions().stream_synchronize(stream.get()); - poisoned = true; - position = 0; - saved.reset(); - history.clear(); - } - throw; - } - } - - buffer Prefill(const std::vector& ids) { - usable(); - if (ids.empty()) throw std::invalid_argument("Qwen3 request contains no token IDs"); - std::vector tokens; - if (position != 0) { - if (history.size() != static_cast(position)) - throw std::runtime_error("Qwen3 cannot continue from position " + std::to_string(position) + - " without the tokens behind it"); - tokens.reserve(history.size() + ids.size()); - tokens = history; - } - tokens.insert(tokens.end(), ids.begin(), ids.end()); - const int resume = position; - position = 0; - try { - auto out = Run(tokens); - history = std::move(tokens); - return out; - } catch (...) { - if (!poisoned) position = resume; - throw; - } - } - - buffer Forward(int id) { - auto out = Run(std::span(&id, 1)); - history.push_back(id); - return out; - } - - void Rewind(int to) { - position = to; - if (history.size() > static_cast(to)) history.resize(static_cast(to)); - } - - buffer ReadCache(bool is_k, int layer, int index) { - usable(); - if (layer < 0 || layer >= c->layers || index < 0 || index >= kMaxSequenceLength) - throw std::out_of_range("Qwen3 cache index is out of range"); - auto lease = runtime->AcquireExecution(); - api->Check(api->functions().stream_synchronize(stream.get()), - "ryzenai_corelib_stream_synchronize cache read"); - const auto heads = static_cast(c->kv_heads); - const auto width = static_cast(c->head_dim); - buffer out(heads * width); - void* cache = is_k ? k_cache[layer].get() : v_cache[layer].get(); - for (std::size_t head = 0; head < heads; ++head) { - const auto offset = (head * kMaxSequenceLength + static_cast(index)) * width; - api->Check(api->functions().tensor_read(cache, ryzenai_corelib_data_type_bf16, - out.data() + head * width, width, offset), - "ryzenai_corelib_tensor_read cache head " + std::to_string(head)); - } - return out; - } -}; - -qwen3_rai::qwen3_rai(LM_Config, std::shared_ptr package, - std::shared_ptr runtime, std::uint32_t max_length) - : impl_(std::make_unique(std::move(package), std::move(runtime), max_length)) {} -qwen3_rai::~qwen3_rai() = default; - -buffer qwen3_rai::forward(int id) { return impl_->Forward(id); } -buffer qwen3_rai::prefill(std::vector& ids, void*) { return impl_->Prefill(ids); } - -void qwen3_rai::set_context_length(int n) { - impl_->usable(); - if (n < 0 || static_cast(n) > impl_->max_length) - throw std::out_of_range("Qwen3 context length is out of range"); - impl_->Rewind(n); -} - -// An ABI shim, not a capability. load_weights is pure virtual in causal_lm.hpp, -// which is frozen because the engine libraries in src/lib/ are prebuilt -// against it. Nothing calls this: this engine's weights come from the GGUF -// package it was constructed with. See AutoModel/model_backend.hpp. -void qwen3_rai::load_weights(Q4NX&) { - impl_->usable(); - throw std::runtime_error("Qwen3 rai weights are loaded only from GGUF"); -} - -void qwen3_rai::update_max_length(std::uint32_t n) { - impl_->usable(); - if (!n || n > kMaxSequenceLength || n < static_cast(impl_->position)) - throw std::out_of_range("Qwen3 maximum length is invalid"); - impl_->max_length = n; -} - -void qwen3_rai::clear_context() { - impl_->usable(); - impl_->position = 0; - impl_->saved.reset(); - impl_->history.clear(); -} - -buffer qwen3_rai::get_k_cache(int l, int i) { return impl_->ReadCache(true, l, i); } -buffer qwen3_rai::get_v_cache(int l, int i) { return impl_->ReadCache(false, l, i); } - -int qwen3_rai::get_current_context_length() { - impl_->usable(); - return impl_->position; -} - -int qwen3_rai::checkpoint() { - impl_->usable(); - impl_->saved = impl_->position; - return impl_->position; -} - -int qwen3_rai::restore() { - impl_->usable(); - if (!impl_->saved) return -1; - impl_->Rewind(*impl_->saved); - return impl_->position; -} - -bool qwen3_rai::poisoned() const noexcept { return impl_ && impl_->poisoned; } - -} // namespace flm::qwen3 diff --git a/src/common/models/qwen3/rai/aie_next/qwen3_rai_backend.cpp b/src/common/models/qwen3/rai/aie_next/qwen3_rai_backend.cpp index b9cc5979f..a562ce2c4 100644 --- a/src/common/models/qwen3/rai/aie_next/qwen3_rai_backend.cpp +++ b/src/common/models/qwen3/rai/aie_next/qwen3_rai_backend.cpp @@ -1,115 +1,21 @@ /// \file qwen3_rai_backend.cpp -/// \brief The ryzenai-corelib backend for Qwen3 +/// \brief Qwen3 0.6B, 1.7B, 4B and 8B on corelib's "qwen3" model #include "models/qwen3/rai/aie_next/qwen3_rai_backend.hpp" -#include "models/qwen3/rai/aie_next/qwen3_rai.hpp" -#include "models/qwen3/rai/aie_next/qwen3_rai_gguf.hpp" -#include "rai/corelib_runtime.hpp" -#include "utils/file_access.hpp" -#include "utils/utils.hpp" - -#include -#include -#include -#include -#include -#include -#include +#include "rai/corelib_model_backend.hpp" namespace flm::qwen3 { -namespace { - -nlohmann::json ReadJson(const std::filesystem::path& path) { - flm::file_access::ObserveOpen(path); - std::ifstream input(path, std::ios::binary); - if (!input) throw std::runtime_error("Cannot open " + path.string()); - try { - return nlohmann::json::parse(input); - } catch (const std::exception& error) { - throw std::runtime_error("Cannot parse " + path.string() + ": " + error.what()); - } -} - -/// \brief the one .gguf the catalog entry lists -/// \note One backend serves four sizes, so the file name comes from the entry -/// rather than from a constant. -std::string GgufName(const nlohmann::ordered_json& model_info) { - std::string found; - if (model_info.contains("files") && model_info["files"].is_array()) { - for (const auto& file : model_info["files"]) { - if (!file.is_string()) continue; - const auto name = file.get(); - if (name.size() > 5 && name.ends_with(".gguf")) { - if (!found.empty()) - throw std::runtime_error("Qwen3 rai entry lists more than one .gguf file"); - found = name; - } - } - } - if (found.empty()) throw std::runtime_error("Qwen3 rai entry lists no .gguf file"); - return found; -} - -class RaiBackend final : public flm::backend::ModelBackend { -public: - explicit RaiBackend(const flm::backend::BackendContext& context) { - if (!context.config) throw std::runtime_error("Qwen3 rai backend needs an LM_Config"); - if (context.enable_preemption) - throw std::invalid_argument("Qwen3 rai does not support preemption"); - if (context.context_length < 1 || context.context_length > kRaiContextLimit) - throw std::out_of_range("Qwen3 rai context length must be in 1..4096"); - - // Validate every source of truth before the runtime or any device - // object exists, so a mismatched package fails with nothing allocated. - const std::filesystem::path root(context.model_path); - const auto config = ReadJson(root / "config.json"); - const auto tokenizer = ReadJson(root / "tokenizer.json"); - const auto tokenizer_config = context.tokenizer_config - ? *context.tokenizer_config - : ReadJson(root / "tokenizer_config.json"); - auto package = Qwen3GgufPackage::Open(root / GgufName(context.model_info)); - package->ValidateContract(config, tokenizer, tokenizer_config); - - runtime_ = corelib::CorelibRuntime::GetOrCreate(utils::get_executable_directory()); - auto engine = std::make_unique(*context.config, std::move(package), runtime_, - context.context_length); - engine->clear_context(); - engine_ = std::move(engine); - } - - ~RaiBackend() override { - // corelib objects must not outlive the API they came from. - engine_.reset(); - } - - causal_lm& engine() override { return *engine_; } - std::string id() const override { return flm::backend::kRaiBackendId; } - std::string detail() const override { - return runtime_ ? runtime_->loaded_library_path().string() : std::string(); - } - std::uint32_t max_decode_length() const override { return kRaiDecodeLimit; } - bool supports_preemption() const override { return false; } - /// \note corelib rejects a decode past its own limit, so the extra forward() - /// the FastFlowLM engines want after an EOS token would fail here. - bool forwards_past_eos() const override { return false; } - bool poisoned() const noexcept override { return engine_ && engine_->poisoned(); } - std::optional> forced_eos_ids() const override { - return std::vector(kEosIds.begin(), kEosIds.end()); - } - -private: - // Declared first so it outlives the engine. - std::shared_ptr runtime_; - std::unique_ptr engine_; -}; - -} // namespace +/// \note A multi-row pass starts only at position 0, so a later prompt re-runs +/// the conversation with it, as the per-op engine did. flm::backend::BackendFactory rai_factory() { - return [](const flm::backend::BackendContext& context) - -> std::unique_ptr { - return std::make_unique(context); - }; + return rai::Factory({ + .model = {"qwen3", "Qwen3", kRaiContextLimit, rai::Continuation::kFromStart, + rai::Rewind::kSeek}, + .context_rule = "1..4096", + // <|im_end|>, <|endoftext|> + .eos_ids = {151645, 151643}, + }); } } // namespace flm::qwen3 diff --git a/src/common/models/qwen3/rai/aie_next/qwen3_rai_config.cpp b/src/common/models/qwen3/rai/aie_next/qwen3_rai_config.cpp deleted file mode 100644 index f61937e94..000000000 --- a/src/common/models/qwen3/rai/aie_next/qwen3_rai_config.cpp +++ /dev/null @@ -1,41 +0,0 @@ -#include "models/qwen3/rai/aie_next/qwen3_rai_config.hpp" - -#include -#include -#include - -namespace flm::qwen3 { -namespace { - -bool Close(double a, double b) { - return std::isfinite(a) && std::fabs(a - b) <= 1e-6 * std::fabs(b); -} - -} // namespace - -const Qwen3Config& SelectQwen3Row(const Qwen3Shape& shape) { - for (const auto& row : kQwen3Rows) { - if (row.layers != shape.layers || row.hidden != shape.hidden || - row.q_heads != shape.q_heads || row.kv_heads != shape.kv_heads || - row.head_dim != shape.head_dim || row.intermediate != shape.intermediate) - continue; - // A GGUF stores these as float32, so 1e-6 reads back as 9.99999997e-07. - if (!Close(shape.epsilon, row.epsilon)) - throw std::runtime_error("Qwen3-" + std::string(row.size) + " expects rms epsilon " + - std::to_string(row.epsilon) + ", but the GGUF says " + - std::to_string(shape.epsilon)); - if (!Close(shape.rope_theta, row.rope_theta)) - throw std::runtime_error("Qwen3-" + std::string(row.size) + " expects rope theta " + - std::to_string(row.rope_theta) + ", but the GGUF says " + - std::to_string(shape.rope_theta)); - return row; - } - std::ostringstream message; - message << "no supported Qwen3 size matches this GGUF: layers=" << shape.layers - << " hidden=" << shape.hidden << " heads=" << shape.q_heads << "/" << shape.kv_heads - << " head_dim=" << shape.head_dim << " intermediate=" << shape.intermediate - << " (the rai backend runs Qwen3 0.6B, 1.7B, 4B and 8B)"; - throw std::runtime_error(message.str()); -} - -} // namespace flm::qwen3 diff --git a/src/common/models/qwen3/rai/aie_next/qwen3_rai_gguf.cpp b/src/common/models/qwen3/rai/aie_next/qwen3_rai_gguf.cpp deleted file mode 100644 index f840b71d2..000000000 --- a/src/common/models/qwen3/rai/aie_next/qwen3_rai_gguf.cpp +++ /dev/null @@ -1,182 +0,0 @@ -#include "models/qwen3/rai/aie_next/qwen3_rai_gguf.hpp" - -#include -#include -#include -#include -#include -#include - -namespace flm::qwen3 { -namespace { -using rai::FailField; - -std::string Name(std::size_t layer, const char* suffix) { - return "blk." + std::to_string(layer) + suffix; -} - -bool IsRetained(std::string_view key) { - return key == "general.architecture" || key.starts_with("qwen3.") || - key == "tokenizer.ggml.eos_token_id"; -} - -std::int64_t SignedMetadata(const rai::GgufFile& file, std::string_view key) { - return static_cast(file.Unsigned(key)); -} - -const nlohmann::json* Find(const nlohmann::json& object, std::string_view key) { - const auto it = object.find(std::string(key)); - return it == object.end() ? nullptr : &*it; -} - -void RequireInteger(const nlohmann::json& object, std::string_view key, std::int64_t expected, - bool optional = false) { - const auto* value = Find(object, key); - if (!value) { - if (optional) return; - FailField(key, "missing", std::to_string(expected)); - } - if (!value->is_number_integer()) FailField(key, value->dump(), std::to_string(expected)); - if (value->get() != expected) - FailField(key, value->dump(), std::to_string(expected)); -} - -void RequireClose(const nlohmann::json& object, std::string_view key, double expected) { - const auto* value = Find(object, key); - if (!value || !value->is_number()) - FailField(key, value ? value->dump() : "missing", std::to_string(expected)); - const double actual = value->get(); - if (!std::isfinite(actual) || std::fabs(actual - expected) > 1e-6 * std::fabs(expected)) - FailField(key, value->dump(), std::to_string(expected)); -} - -void RequireString(const nlohmann::json& object, std::string_view key, std::string_view expected) { - const auto* value = Find(object, key); - if (!value || !value->is_string() || value->get_ref() != expected) - FailField(key, value ? value->dump() : "missing", std::string(expected)); -} -} // namespace - -Qwen3GgufPackage::~Qwen3GgufPackage() = default; - -std::shared_ptr Qwen3GgufPackage::Open(const std::filesystem::path& gguf_path) { - std::shared_ptr package(new Qwen3GgufPackage()); - package->file_ = rai::GgufFile::Open(gguf_path, IsRetained); - const auto& file = *package->file_; - - // The architecture string, not a missing tensor, is what says QK-Norm is - // there: a Qwen2 conversion simply has no attn_q_norm to find. - const auto architecture = file.String("general.architecture"); - if (architecture != "qwen3") FailField("general.architecture", architecture, "qwen3"); - Qwen3Shape shape{}; - shape.layers = SignedMetadata(file, "qwen3.block_count"); - shape.hidden = SignedMetadata(file, "qwen3.embedding_length"); - shape.q_heads = SignedMetadata(file, "qwen3.attention.head_count"); - shape.kv_heads = SignedMetadata(file, "qwen3.attention.head_count_kv"); - // Required, not derived: 0.6B is 16 heads of 128 over a hidden size of - // 1024, where hidden / heads would give 64. - shape.head_dim = SignedMetadata(file, "qwen3.attention.key_length"); - if (file.HasMetadata("qwen3.attention.value_length") && - SignedMetadata(file, "qwen3.attention.value_length") != shape.head_dim) - FailField("qwen3.attention.value_length", - std::to_string(file.Unsigned("qwen3.attention.value_length")), - std::to_string(shape.head_dim)); - shape.intermediate = SignedMetadata(file, "qwen3.feed_forward_length"); - shape.rope_theta = file.Number("qwen3.rope.freq_base"); - shape.epsilon = file.Number("qwen3.attention.layer_norm_rms_epsilon"); - package->config_ = &SelectQwen3Row(shape); - const auto& c = *package->config_; - - if (file.Unsigned("tokenizer.ggml.eos_token_id") != static_cast(kEosIds[0])) - FailField("tokenizer.ggml.eos_token_id", - std::to_string(file.Unsigned("tokenizer.ggml.eos_token_id")), - std::to_string(kEosIds[0])); - - const std::array hidden{c.hidden}; - const std::array head{c.head_dim}; - const std::array table{c.vocab, c.hidden}; - package->embedding_ = file.RequireQ8("token_embd.weight", table); - package->lm_head_ = file.HasTensor("output.weight") ? file.RequireQ8("output.weight", table) - : package->embedding_; - package->output_norm_ = file.RequireF32("output_norm.weight", hidden); - package->layers_.reserve(static_cast(c.layers)); - for (std::size_t i = 0; i < static_cast(c.layers); ++i) { - Qwen3LayerTensors t; - t.attn_norm = file.RequireF32(Name(i, ".attn_norm.weight"), hidden); - t.ffn_norm = file.RequireF32(Name(i, ".ffn_norm.weight"), hidden); - t.q_norm = file.RequireF32(Name(i, ".attn_q_norm.weight"), head); - t.k_norm = file.RequireF32(Name(i, ".attn_k_norm.weight"), head); - // [out, in]: Q is hidden -> q_heads * head_dim, which is NOT hidden on - // 0.6B (2048 vs 1024) or 4B (4096 vs 2560); O is the reverse. - t.q = file.RequireQ8(Name(i, ".attn_q.weight"), std::array{c.q_dim(), c.hidden}); - t.k = file.RequireQ8(Name(i, ".attn_k.weight"), std::array{c.kv_dim(), c.hidden}); - t.v = file.RequireQ8(Name(i, ".attn_v.weight"), std::array{c.kv_dim(), c.hidden}); - t.o = file.RequireQ8(Name(i, ".attn_output.weight"), std::array{c.hidden, c.q_dim()}); - t.gate = file.RequireQ8(Name(i, ".ffn_gate.weight"), std::array{c.intermediate, c.hidden}); - t.up = file.RequireQ8(Name(i, ".ffn_up.weight"), std::array{c.intermediate, c.hidden}); - t.down = file.RequireQ8(Name(i, ".ffn_down.weight"), std::array{c.hidden, c.intermediate}); - package->layers_.push_back(std::move(t)); - } - return package; -} - -const Qwen3Config& Qwen3GgufPackage::Config() const noexcept { return *config_; } -const std::filesystem::path& Qwen3GgufPackage::Path() const { return file_->Path(); } -const rai::GgufTensor& Qwen3GgufPackage::Embedding() const noexcept { return embedding_; } -const rai::GgufTensor& Qwen3GgufPackage::LmHead() const noexcept { return lm_head_; } -bool Qwen3GgufPackage::TiedLmHead() const noexcept { return lm_head_.name == embedding_.name; } -const rai::GgufFloatTensor& Qwen3GgufPackage::OutputNorm() const noexcept { return output_norm_; } - -const Qwen3LayerTensors& Qwen3GgufPackage::Layer(std::size_t layer) const { - if (layer >= layers_.size()) - FailField("Qwen3 layer", std::to_string(layer), "0.." + std::to_string(layers_.size() - 1)); - return layers_[layer]; -} - -void Qwen3GgufPackage::ValidateContract(const nlohmann::json& config, - const nlohmann::json& tokenizer, - const nlohmann::json& tokenizer_config) const { - const auto& c = *config_; - RequireString(config, "model_type", "qwen3"); - RequireInteger(config, "num_hidden_layers", c.layers); - RequireInteger(config, "hidden_size", c.hidden); - RequireInteger(config, "intermediate_size", c.intermediate); - RequireInteger(config, "num_attention_heads", c.q_heads); - RequireInteger(config, "num_key_value_heads", c.kv_heads); - RequireInteger(config, "head_dim", c.head_dim, true); - RequireInteger(config, "vocab_size", c.vocab); - RequireClose(config, "rms_norm_eps", c.epsilon); - RequireClose(config, "rope_theta", c.rope_theta); - if (const auto* tied = Find(config, "tie_word_embeddings")) { - if (!tied->is_boolean() || tied->get() != TiedLmHead()) - FailField("tie_word_embeddings", tied->dump(), - TiedLmHead() ? "true (the GGUF has no output.weight)" - : "false (the GGUF ships output.weight)"); - } - - std::map> special; - try { - for (const auto& item : tokenizer.at("added_tokens")) - special.emplace(item.at("content").get(), item.at("id").get()); - } catch (const nlohmann::json::exception& error) { - FailField("tokenizer.json added_tokens", error.what(), "an array of {content, id}"); - } - for (const auto& [token, id] : std::array, 5>{{ - {"<|endoftext|>", kEosIds[1]}, {"<|im_start|>", 151644}, {"<|im_end|>", kEosIds[0]}, - {"", 151667}, {"", 151668}}}) { - const auto it = special.find(token); - if (it == special.end()) FailField(token, "missing from tokenizer.json", std::to_string(id)); - if (it->second != id) FailField(token, std::to_string(it->second), std::to_string(id)); - } - - RequireString(tokenizer_config, "eos_token", "<|im_end|>"); - const auto* chat_template = Find(tokenizer_config, "chat_template"); - if (!chat_template || !chat_template->is_string()) - FailField("chat_template", chat_template ? chat_template->dump() : "missing", - "string containing the ChatML markers"); - for (const auto marker : {"<|im_start|>", "<|im_end|>"}) - if (chat_template->get_ref().find(marker) == std::string::npos) - FailField(marker, "missing from chat_template", "present in chat_template"); -} - -} // namespace flm::qwen3 diff --git a/src/common/models/qwen3/rai/aie_next/qwen3_rai_shape_plan.cpp b/src/common/models/qwen3/rai/aie_next/qwen3_rai_shape_plan.cpp deleted file mode 100644 index cca288bea..000000000 --- a/src/common/models/qwen3/rai/aie_next/qwen3_rai_shape_plan.cpp +++ /dev/null @@ -1,83 +0,0 @@ -#include "models/qwen3/rai/aie_next/qwen3_rai_shape_plan.hpp" - -#include "rai/kernel_grid.hpp" - -#include -#include -#include - -namespace flm::qwen3 { -namespace { - -std::string Shape(std::int64_t m, std::int64_t k, std::int64_t n) { - return "[" + std::to_string(m) + "," + std::to_string(k) + "]x[" + std::to_string(k) + "," + - std::to_string(n) + "]"; -} - -void RequireMatmul(const corelib::ShapeGrid& grid, std::int64_t rows, std::int64_t k, - std::int64_t n, std::uint32_t group, const char* name) { - const auto call = std::string(name) + " " + Shape(rows, k, n) + " g" + std::to_string(group); - const auto covered = grid.SmallestM(rows, k, n, group); - if (covered != rows) - throw std::runtime_error(call + ": corelib does not ship this bucket exactly" + - (covered < 0 ? "" : " (nearest " + std::to_string(covered) + ")")); -} - -void RequireQkNormRows(std::int64_t rows, std::int64_t heads, const char* which) { - const auto norm_rows = rows * heads; - if (std::find(kQkNormRows.begin(), kQkNormRows.end(), norm_rows) == kQkNormRows.end()) - throw std::runtime_error(std::string(which) + "-norm would run " + std::to_string(rows) + - " rows x " + std::to_string(heads) + " heads = " + - std::to_string(norm_rows) + - " rows at k=128, which no RMSNorm kernel ships"); -} - -} // namespace - -Qwen3ShapePlan Qwen3ShapePlan::Build(const std::shared_ptr& api, - ryzenai_corelib_stream_ptr stream, - const Qwen3Config& c) { - if (!api) throw std::invalid_argument("Qwen3ShapePlan corelib API is null"); - if (!stream) throw std::invalid_argument("Qwen3ShapePlan stream is null"); - if (c.head_dim != kQkNormWidth) - throw std::runtime_error("Qwen3 head_dim " + std::to_string(c.head_dim) + - " is not the width the QK-Norm grid was read at"); - - Qwen3ShapePlan plan; - // Full rotary: rope_dim == head_dim. A partial one selects a different - // attention ELF rather than configuring the same one. - plan.attention_ = {c.q_heads, c.kv_heads, c.head_dim, kMaxSequenceLength, c.head_dim}; - const auto matmul = corelib::ShapeGrid::Matmul(*api, stream); - const auto ssmlp = corelib::ShapeGrid::SsMlp(*api, stream, false); - const auto mha = corelib::ShapeGrid::FlatMha(*api, stream, plan.attention_); - - for (const auto rows : kExecutionRows) { - RequireMatmul(matmul, rows, c.hidden, c.q_dim(), c.group, "query"); - RequireMatmul(matmul, rows, c.hidden, c.kv_dim(), c.group, "key/value"); - RequireMatmul(matmul, rows, c.q_dim(), c.hidden, c.group, "output"); - - const auto ssmlp_rows = ssmlp.SmallestM(rows, c.hidden, c.intermediate, c.group); - if (ssmlp_rows != rows) - throw std::runtime_error("ssmlp [" + std::to_string(rows) + "," + std::to_string(c.hidden) + - "," + std::to_string(c.intermediate) + - "] does not ship exactly"); - - const auto mha_rows = mha.SmallestM(rows, 0, 0, -1); - if (mha_rows != rows) - throw std::runtime_error("flat_mha [" + std::to_string(rows) + - "] does not ship exactly"); - - RequireQkNormRows(rows, c.q_heads, "q"); - RequireQkNormRows(rows, c.kv_heads, "k"); - } - RequireMatmul(matmul, 1, c.hidden, c.vocab, c.head_group, "lm_head"); - return plan; -} - -std::int64_t Qwen3ShapePlan::RowsFor(std::size_t live_rows) const { - for (const auto rows : kExecutionRows) - if (live_rows > 0 && static_cast(live_rows) <= rows) return rows; - throw std::out_of_range("Qwen3 live rows must be in 1.." + std::to_string(kMaxSequenceLength)); -} - -} // namespace flm::qwen3 diff --git a/src/common/models/qwen35/rai/aie_next/qwen35_rai.cpp b/src/common/models/qwen35/rai/aie_next/qwen35_rai.cpp deleted file mode 100644 index ba413319b..000000000 --- a/src/common/models/qwen35/rai/aie_next/qwen35_rai.cpp +++ /dev/null @@ -1,746 +0,0 @@ -#include "models/qwen35/rai/aie_next/qwen35_rai.hpp" - -#include "rai/corelib_object.hpp" -#include "rai/host_ops.hpp" -#include "rai/kernel_grid.hpp" -#include "rai/weight_source.hpp" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace flm::qwen35 { -namespace { -using namespace flm::corelib; - -int PadVectors(int vectors, int width) { - const int* table = nullptr; - int count = 0; - if (width == 128) { - static const int kRows[] = {2, 16, 32}; - table = kRows; - count = 3; - } else if (width == 256) { - static const int kRows[] = {1, 8, 32}; - table = kRows; - count = 3; - } else { - throw std::runtime_error("no RMSNorm vector table for width " + std::to_string(width)); - } - for (int i = 0; i < count; ++i) - if (table[i] >= vectors) return table[i]; - throw std::runtime_error("no RMSNorm kernel covers " + std::to_string(vectors) + - " vectors of " + std::to_string(width)); -} - -std::uint16_t FloatToHalfOne(float value) { - const auto bits = std::bit_cast(value); - const std::uint32_t sign = (bits >> 16) & 0x8000u; - auto exponent = static_cast((bits >> 23) & 0xffu) - 127 + 15; - std::uint32_t mantissa = bits & 0x7fffffu; - if (exponent <= 0) { - if (exponent < -10) return static_cast(sign); - mantissa |= 0x800000u; - const auto shift = static_cast(14 - exponent); - std::uint32_t half = mantissa >> shift; - const auto remainder = mantissa & ((1u << shift) - 1u); - if (remainder > (1u << (shift - 1)) || - (remainder == (1u << (shift - 1)) && (half & 1u))) - ++half; - return static_cast(sign | half); - } - if (exponent >= 31) return static_cast(sign | 0x7c00u); - std::uint32_t half = mantissa >> 13; - const auto remainder = mantissa & 0x1fffu; - if (remainder > 0x1000u || (remainder == 0x1000u && (half & 1u))) { - ++half; - if (half == 0x400u) { - half = 0; - ++exponent; - if (exponent >= 31) return static_cast(sign | 0x7c00u); - } - } - return static_cast(sign | (static_cast(exponent) << 10) | half); -} - -std::vector FloatToHalf(const std::vector& values) { - std::vector out(values.size()); - for (std::size_t i = 0; i < values.size(); ++i) out[i] = FloatToHalfOne(values[i]); - return out; -} - -struct FullWeights { - UniqueMatMulWeights q, gate, k, v, o; - UniqueRmsNormWeights q_norm, k_norm; - UniqueSsMlpWeights mlp; -}; - -struct LinearWeights { - UniqueMatMulWeights qkv, beta, alpha, z, out; - UniqueTensor neg_a; - UniqueDwConvWeights conv; - UniqueRmsNormWeights norm; - UniqueSsMlpWeights mlp; -}; - -struct FullIo { - UniqueTensor q_raw, q_buf, k_buf, mha_out, gated, k_cache, v_cache; - UniqueTensor q_src[2], q_dst[2], q_live, k_live; - std::vector v_rows; -}; - -struct LinearIo { - UniqueTensor qkv, qkv_row, beta, beta_live, g, g_live; - UniqueTensor dw_q, dw_k, dw_v, lin_out, lin_flat, z_out; - UniqueTensor conv[2], state[2]; -}; - -} // namespace - -using namespace flm::corelib; - -class qwen35_rai::Impl { -public: - Impl(std::shared_ptr package, std::shared_ptr runtime, - std::uint32_t maximum) - : package(std::move(package)), runtime(std::move(runtime)), - api(this->runtime ? this->runtime->api() : nullptr), max_length(maximum) { - if (!this->package) throw std::invalid_argument("Qwen3.5 GGUF package is null"); - if (!this->runtime || !api) throw std::invalid_argument("corelib runtime is null"); - if (!maximum || maximum > kMaxSequenceLength) - throw std::invalid_argument("Qwen3.5 maximum length must be in 1..4096"); - cfg = &this->package->Config(); - const int layers = cfg->layers; - - auto lease = this->runtime->AcquireExecution(); - stream = corelib::OpenModelStream(api, kPrefillPdi, kTokenPdi, "qwen35"); - - const auto matmul = ShapeGrid::Matmul(*api, stream.get()); - const auto sigmoid = ShapeGrid::MatmulAct( - *api, stream.get(), ryzenai_corelib_matmul_activation_sigmoid); - const auto softplus = ShapeGrid::MatmulAct( - *api, stream.get(), ryzenai_corelib_matmul_activation_softplus_mulbcast); - const auto silu = ShapeGrid::MatmulAct( - *api, stream.get(), ryzenai_corelib_matmul_activation_silu_mul); - const auto sigmoid_mul = ShapeGrid::MatmulAct( - *api, stream.get(), ryzenai_corelib_matmul_activation_sigmoid_mul); - const auto mlp_grid = ShapeGrid::SsMlp(*api, stream.get(), false); - const auto require = [](const ShapeGrid& grid, int k, int n, const char* name) { - if (!grid.Has(1, k, n, -1)) - throw std::runtime_error(std::string("no ") + name + " kernel ships at 1x" + - std::to_string(k) + "x" + std::to_string(n)); - }; - require(matmul, cfg->hidden, cfg->q_dim(), "matmul"); - require(matmul, cfg->hidden, cfg->kv_dim(), "matmul"); - require(matmul, cfg->q_dim(), cfg->hidden, "matmul"); - require(matmul, cfg->hidden, cfg->conv_channels(), "matmul"); - require(matmul, cfg->lin_value_dim(), cfg->hidden, "matmul"); - require(matmul, cfg->hidden, cfg->vocab, "matmul"); - require(sigmoid, cfg->hidden, cfg->lin_v_heads, "sigmoid"); - require(softplus, cfg->hidden, cfg->lin_v_heads, "softplus_mulbcast"); - require(silu, cfg->hidden, cfg->lin_value_dim(), "silu_mul"); - require(sigmoid_mul, cfg->hidden, cfg->q_dim(), "sigmoid_mul"); - require(mlp_grid, cfg->hidden, cfg->intermediate, "fused MLP"); - beta_width = static_cast(sigmoid.SmallestN(1, cfg->hidden, cfg->lin_v_heads)); - alpha_width = static_cast(softplus.SmallestN(1, cfg->hidden, cfg->lin_v_heads)); - if (beta_width < 0 || alpha_width < 0) - throw std::runtime_error("no gate kernel is wide enough for the linear-attention heads"); - - cfg_eps = static_cast(cfg->eps); - epsilon_bf16 = rai::F32ToBf16(std::span(&cfg_eps, 1)); - lin_desc = {cfg->lin_k_heads, cfg->lin_v_heads, cfg->lin_k_dim, cfg->lin_v_dim}; - mha_desc = {cfg->q_heads, cfg->kv_heads, cfg->head_dim, kMaxSequenceLength, cfg->rope_dim, - 0, 0, 0.0f}; - const auto started = std::chrono::steady_clock::now(); - std::cout << "[FLM] Packing Qwen3.5 hidden " << cfg->hidden << ", " << layers - << " layers, from Q8_0" << std::endl; - full.resize(static_cast(layers)); - linear.resize(static_cast(layers)); - norm0 = Norm("blk.0.attn_norm.weight", cfg->hidden); - for (int layer = 0; layer < layers; ++layer) LoadLayer(layer); - lm_head = Matmul(this->package->Blocks(this->package->LmHeadName(), cfg->hidden, cfg->vocab), - cfg->hidden, cfg->vocab, "lm_head"); - const auto seconds = std::chrono::duration(std::chrono::steady_clock::now() - started).count(); - std::cout << "[FLM] Qwen3.5 weights ready in " << seconds << "s" << std::endl; - - Allocate(); - Reset(); - // A process's first M=1 decode is perturbed by a couple of bf16 ulps. - // A throwaway token keeps that pass out of the caller's sequence. - StepBody(0); - Reset(); - } - - buffer Forward(int id) { - usable(); - auto lease = runtime->AcquireExecution(); - auto out = StepBody(id); - history.push_back(id); - return out; - } - - buffer Prefill(const std::vector& ids) { - usable(); - if (ids.empty()) throw std::invalid_argument("Qwen3.5 request contains no token IDs"); - auto lease = runtime->AcquireExecution(); - buffer out; - for (const int id : ids) { - out = StepBody(id); - history.push_back(id); - } - return out; - } - - void Clear() { - usable(); - auto lease = runtime->AcquireExecution(); - Reset(); - history.clear(); - saved.reset(); - } - - void Rewind(int to) { - usable(); - if (to < 0 || static_cast(to) > max_length) - throw std::out_of_range("Qwen3.5 context length is out of range"); - if (static_cast(to) > history.size()) - throw std::runtime_error("Qwen3.5 cannot rewind past the tokens it has run"); - if (to == position) return; - auto lease = runtime->AcquireExecution(); - Reset(); - for (int i = 0; i < to; ++i) StepBody(history[static_cast(i)]); - } - - buffer ReadCache(bool is_k, int layer, int index) { - usable(); - if (layer < 0 || layer >= cfg->layers || !cfg->is_full(layer) || index < 0 || - index >= kMaxSequenceLength) - throw std::out_of_range("Qwen3.5 cache index is out of range"); - auto lease = runtime->AcquireExecution(); - api->Check(api->functions().stream_synchronize(stream.get()), - "ryzenai_corelib_stream_synchronize cache read"); - const auto heads = static_cast(cfg->kv_heads); - const auto width = static_cast(cfg->head_dim); - buffer out(heads * width); - void* cache = is_k ? full_io[layer].k_cache.get() : full_io[layer].v_cache.get(); - for (std::size_t head = 0; head < heads; ++head) { - const auto offset = (head * kMaxSequenceLength + static_cast(index)) * width; - api->Check(api->functions().tensor_read(cache, ryzenai_corelib_data_type_bf16, - out.data() + head * width, width, offset), - "ryzenai_corelib_tensor_read cache"); - } - return out; - } - - void usable() const { - if (poisoned) throw std::runtime_error("Qwen3.5 corelib engine is poisoned"); - } - - std::uint32_t max_length; - int position = 0; - std::optional saved; - bool poisoned = false; - -private: - UniqueTensor Tensor(std::initializer_list shape, const char* label) { - std::vector dims(shape); - void* created = nullptr; - api->Check(api->functions().create_device_tensor(ryzenai_corelib_data_type_bf16, dims.data(), - dims.size(), &created), - label); - return UniqueTensor(api, created); - } - - UniqueTensor Window(void* parent, std::initializer_list shape, std::size_t offset, - const char* label, std::initializer_list strides = {}) { - std::vector dims(shape); - std::vector steps(strides); - void* created = nullptr; - api->Check(api->functions().create_tensor_window(parent, dims.data(), dims.size(), offset, - steps.empty() ? nullptr : steps.data(), - steps.size(), &created), - label); - return UniqueTensor(api, created); - } - - void Write(void* tensor, const float* data, std::size_t count, const char* label) { - api->Check(api->functions().tensor_write(tensor, ryzenai_corelib_data_type_fp32, data, count, 0), - label); - } - - void Zero(void* tensor, std::size_t count, const char* label) { - if (zeros.size() < count) zeros.assign(count, 0.0f); - Write(tensor, zeros.data(), count, label); - } - - UniqueMatMulWeights Matmul(std::span blocks, int k, int n, const std::string& label) { - const ryzenai_corelib_matmul_bf16_weights_desc desc{k, n, kGroup, false}; - const ryzenai_corelib_weights_source sources[]{corelib::GgufQ8( - ryzenai_corelib_weights_role_qweight, blocks.data(), blocks.size(), n, k)}; - const auto options = corelib::PackOptions(kRequantizeThreads); - void* created = nullptr; - api->Check(api->functions().matmul_weights_create( - &desc, sources, std::size(sources), &options, &created), - "ryzenai_corelib_matmul_bf16_weights_create " + label); - return UniqueMatMulWeights(api, created); - } - - UniqueMatMulWeights MatmulOwned(std::vector blocks, int k, int n, const std::string& label) { - return Matmul(blocks, k, n, label); - } - - UniqueMatMulWeights MatmulBiased(std::vector blocks, int k, int n, - const std::vector& bias, const std::string& label) { - const auto onnx = Q8ToOnnx(blocks, k, n, kGroup); - const ryzenai_corelib_matmul_bf16_weights_desc desc{k, n, kGroup, true}; - const auto groups = k / kGroup; - const ryzenai_corelib_weights_source sources[]{ - corelib::MemorySource(ryzenai_corelib_weights_role_qweight, onnx.qweight.data(), - onnx.qweight.size(), - ryzenai_corelib_weights_data_type_onnx_uint8, {n, k / 2}), - corelib::Fp16(ryzenai_corelib_weights_role_scales, onnx.scales.data(), - onnx.scales.size(), {n, groups}), - corelib::MemorySource(ryzenai_corelib_weights_role_qzeros, onnx.qzeros.data(), - onnx.qzeros.size(), - ryzenai_corelib_weights_data_type_onnx_uint8, {n, groups / 2}), - corelib::Fp32(ryzenai_corelib_weights_role_bias, bias.data(), bias.size(), {n})}; - const auto options = corelib::PackOptions(kRequantizeThreads); - void* created = nullptr; - api->Check(api->functions().matmul_weights_create(&desc, sources, std::size(sources), - &options, &created), - "ryzenai_corelib_matmul_bf16_weights_create " + label); - return UniqueMatMulWeights(api, created); - } - - UniqueRmsNormWeights Norm(std::string_view name, int k) { - const std::int64_t shape[] = {k}; - const auto scale = rai::F32ToBf16(package->F32(name, shape)); - const ryzenai_corelib_rmsnorm_bf16_weights_desc desc{k, static_cast(cfg->eps)}; - const ryzenai_corelib_weights_source sources[]{ - corelib::Bf16(ryzenai_corelib_weights_role_scale, scale.data(), scale.size())}; - void* created = nullptr; - api->Check(api->functions().rmsnorm_weights_create(&desc, sources, std::size(sources), - nullptr, &created), - "ryzenai_corelib_rmsnorm_bf16_weights_create " + std::string(name)); - return UniqueRmsNormWeights(api, created); - } - - UniqueRmsNormWeights NormValues(const std::vector& values, int k, const std::string& label) { - const auto scale = rai::F32ToBf16(values); - const ryzenai_corelib_rmsnorm_bf16_weights_desc desc{k, static_cast(cfg->eps)}; - const ryzenai_corelib_weights_source sources[]{ - corelib::Bf16(ryzenai_corelib_weights_role_scale, scale.data(), scale.size())}; - void* created = nullptr; - api->Check(api->functions().rmsnorm_weights_create(&desc, sources, std::size(sources), - nullptr, &created), - "ryzenai_corelib_rmsnorm_bf16_weights_create " + label); - return UniqueRmsNormWeights(api, created); - } - - UniqueSsMlpWeights Mlp(int layer) { - const auto gate = package->Blocks("blk." + std::to_string(layer) + ".ffn_gate", cfg->hidden, - cfg->intermediate); - const auto up = package->Blocks("blk." + std::to_string(layer) + ".ffn_up", cfg->hidden, - cfg->intermediate); - const auto down = package->Blocks("blk." + std::to_string(layer) + ".ffn_down", cfg->intermediate, - cfg->hidden); - const std::int64_t norm_shape[] = {cfg->hidden}; - const auto norm0_scale = rai::F32ToBf16(package->F32( - "blk." + std::to_string(layer) + ".post_attention_norm.weight", norm_shape)); - const auto next_name = layer + 1 == cfg->layers - ? std::string("output_norm.weight") - : "blk." + std::to_string(layer + 1) + ".attn_norm.weight"; - const auto norm1_scale = rai::F32ToBf16(package->F32(next_name, norm_shape)); - const ryzenai_corelib_ssmlp_bf16_weights_desc desc{cfg->hidden, cfg->intermediate, kGroup, 0, 0}; - const ryzenai_corelib_weights_source sources[]{ - corelib::Bf16(ryzenai_corelib_weights_role_epsilon, epsilon_bf16.data(), epsilon_bf16.size()), - corelib::Bf16(ryzenai_corelib_weights_role_norm0, norm0_scale.data(), norm0_scale.size()), - corelib::Bf16(ryzenai_corelib_weights_role_norm1, norm1_scale.data(), norm1_scale.size()), - corelib::GgufQ8(ryzenai_corelib_weights_role_gate_qweight, gate.data(), gate.size(), cfg->intermediate, cfg->hidden), - corelib::GgufQ8(ryzenai_corelib_weights_role_up_qweight, up.data(), up.size(), cfg->intermediate, cfg->hidden), - corelib::GgufQ8(ryzenai_corelib_weights_role_down_qweight, down.data(), down.size(), cfg->hidden, cfg->intermediate)}; - const auto options = corelib::PackOptions(kRequantizeThreads); - void* created = nullptr; - api->Check(api->functions().ssmlp_weights_create( - &desc, sources, std::size(sources), &options, &created), - "ryzenai_corelib_ssmlp_bf16_weights_create layer " + - std::to_string(layer)); - return UniqueSsMlpWeights(api, created); - } - - UniqueMatMulWeights SliceHeads(UniqueMatMulWeights& both, bool gate_half) { - const int d = cfg->head_dim; - std::vector ranges(static_cast(cfg->q_heads)); - for (int head = 0; head < cfg->q_heads; ++head) { - const std::int64_t begin = gate_half ? (2 * d * head + d) : (2 * d * head); - ranges[static_cast(head)] = {1, begin, begin + d}; - } - void* created = nullptr; - api->Check(api->functions().weights_slice(both.get(), ranges.data(), ranges.size(), &created), - "ryzenai_corelib_weights_slice attn_q"); - return UniqueMatMulWeights(api, created); - } - - void LoadLayer(int layer) { - const auto id = std::to_string(layer); - if (cfg->is_full(layer)) { - auto both = Matmul(package->Blocks("blk." + id + ".attn_q", cfg->hidden, 2 * cfg->q_dim()), - cfg->hidden, 2 * cfg->q_dim(), "blk." + id + ".attn_q"); - full[layer].q = SliceHeads(both, false); - full[layer].gate = SliceHeads(both, true); - both.reset(); - full[layer].k = Matmul(package->Blocks("blk." + id + ".attn_k", cfg->hidden, cfg->kv_dim()), - cfg->hidden, cfg->kv_dim(), "blk." + id + ".attn_k"); - full[layer].v = Matmul(package->Blocks("blk." + id + ".attn_v", cfg->hidden, cfg->kv_dim()), - cfg->hidden, cfg->kv_dim(), "blk." + id + ".attn_v"); - full[layer].o = Matmul(package->Blocks("blk." + id + ".attn_output", cfg->q_dim(), cfg->hidden), - cfg->q_dim(), cfg->hidden, "blk." + id + ".attn_output"); - full[layer].q_norm = Norm("blk." + id + ".attn_q_norm.weight", cfg->head_dim); - full[layer].k_norm = Norm("blk." + id + ".attn_k_norm.weight", cfg->head_dim); - full[layer].mlp = Mlp(layer); - } else { - auto neg = package->PerHead(layer, "ssm_a"); - std::vector operand(static_cast(kSoftplusOperand), 0.0f); - for (int head = 0; head < cfg->lin_v_heads && head < static_cast(neg.size()); ++head) - operand[static_cast(head)] = neg[static_cast(head)]; - linear[layer].neg_a = Tensor({kSoftplusOperand}, "neg_a"); - Write(linear[layer].neg_a.get(), operand.data(), operand.size(), "neg_a"); - linear[layer].qkv = MatmulOwned(package->QkvBlocks(layer), cfg->hidden, cfg->conv_channels(), - "blk." + id + ".attn_qkv"); - linear[layer].beta = MatmulOwned(package->GateBlocks(layer, "beta"), cfg->hidden, cfg->lin_v_heads, - "blk." + id + ".ssm_beta"); - linear[layer].alpha = MatmulBiased(package->GateBlocks(layer, "alpha"), cfg->hidden, - cfg->lin_v_heads, package->PerHead(layer, "ssm_dt.bias"), - "blk." + id + ".ssm_alpha"); - const auto taps = package->ConvTaps(layer); - const auto weights = FloatToHalf(taps); - std::vector bias(static_cast(cfg->conv_channels()), 0); - const ryzenai_corelib_dwconv_bf16_weights_desc desc{ - cfg->conv_channels(), cfg->conv_width, cfg->lin_k_dim, cfg->lin_k_heads, cfg->lin_k_heads, - cfg->lin_v_heads, 1.0e-6f, static_cast(1.0 / std::sqrt(cfg->lin_k_dim)), 1.0f}; - const ryzenai_corelib_weights_source sources[]{ - corelib::Fp16(ryzenai_corelib_weights_role_weights, weights.data(), weights.size(), - {desc.channels, 1, desc.kernel_width}), - corelib::Fp16(ryzenai_corelib_weights_role_bias, bias.data(), bias.size(), - {desc.channels})}; - void* created = nullptr; - api->Check(api->functions().dwconv_weights_create(&desc, sources, std::size(sources), - nullptr, &created), - "ryzenai_corelib_dwconv_bf16_weights_create blk." + id); - linear[layer].conv = UniqueDwConvWeights(api, created); - linear[layer].z = MatmulOwned(package->ZBlocks(layer), cfg->hidden, cfg->lin_value_dim(), - "blk." + id + ".attn_gate"); - linear[layer].norm = Norm("blk." + id + ".ssm_norm.weight", cfg->lin_v_dim); - linear[layer].out = MatmulOwned(package->OutBlocks(layer), cfg->lin_value_dim(), cfg->hidden, - "blk." + id + ".ssm_out"); - linear[layer].mlp = Mlp(layer); - } - std::cout << "[FLM] layer " << (layer + 1) << "/" << cfg->layers << " packed (" - << (cfg->is_full(layer) ? "full" : "linear") << ")" << std::endl; - } - - void Allocate() { - const int layers = cfg->layers; - const int half = cfg->q_heads / 2; - if (cfg->q_heads % 2 != 0 || PadVectors(half, cfg->head_dim) != half) - throw std::runtime_error("q_norm over " + std::to_string(cfg->q_heads) + - " heads does not split into two shipped vector counts"); - q_vectors = std::max(PadVectors(cfg->q_heads, cfg->head_dim), cfg->q_heads); - k_vectors = std::max(PadVectors(cfg->kv_heads, cfg->head_dim), cfg->kv_heads); - input_rows = std::max(1, cfg->conv_width - 1); - BuildRope(); - - res_in = Tensor({1, cfg->hidden}, "residual"); - h_in = Tensor({1, cfg->hidden}, "hidden"); - logits = Tensor({1, cfg->vocab}, "logits"); - res.reserve(layers); - h.reserve(layers); - attn.reserve(layers); - for (int layer = 0; layer < layers; ++layer) { - res.push_back(Tensor({1, cfg->hidden}, "residual")); - h.push_back(Tensor({1, cfg->hidden}, "hidden")); - attn.push_back(Tensor({1, cfg->hidden}, "attention")); - } - full_io = std::vector(static_cast(layers)); - linear_io = std::vector(static_cast(layers)); - parity.assign(static_cast(layers), 0); - for (int layer = 0; layer < layers; ++layer) { - if (!cfg->is_full(layer)) { - auto& io = linear_io[layer]; - io.qkv = Tensor({input_rows, cfg->conv_channels()}, "qkv"); - io.qkv_row = Window(io.qkv.get(), {1, cfg->conv_channels()}, 0, "qkv row"); - io.beta = Tensor({1, beta_width}, "beta"); - io.beta_live = Window(io.beta.get(), {1, cfg->lin_v_heads}, 0, "beta live"); - io.g = Tensor({1, alpha_width}, "gate"); - io.g_live = Window(io.g.get(), {1, cfg->lin_v_heads}, 0, "gate live"); - io.dw_q = Tensor({cfg->lin_k_heads, 1, cfg->lin_k_dim}, "dw q"); - io.dw_k = Tensor({cfg->lin_k_heads, 1, cfg->lin_k_dim}, "dw k"); - io.dw_v = Tensor({cfg->lin_v_heads, 1, cfg->lin_v_dim}, "dw v"); - io.lin_out = Tensor({1, cfg->lin_v_heads, cfg->lin_v_dim}, "linear out"); - io.lin_flat = Window(io.lin_out.get(), {1, cfg->lin_value_dim()}, 0, "linear flat"); - io.z_out = Tensor({1, cfg->lin_value_dim()}, "z"); - for (int side = 0; side < 2; ++side) { - io.conv[side] = Tensor({input_rows, cfg->conv_channels()}, "conv state"); - io.state[side] = Tensor({cfg->lin_v_heads, cfg->lin_v_dim, cfg->lin_k_dim}, "recurrent state"); - } - continue; - } - auto& io = full_io[layer]; - io.q_raw = Tensor({1, cfg->q_dim()}, "q"); - io.q_buf = Tensor({1, q_vectors, cfg->head_dim}, "q buffer"); - io.k_buf = Tensor({1, k_vectors, cfg->head_dim}, "k buffer"); - for (int part = 0; part < 2; ++part) { - const auto offset = static_cast(part * half * cfg->head_dim); - io.q_src[part] = Window(io.q_raw.get(), {1, half, cfg->head_dim}, offset, "q source"); - io.q_dst[part] = Window(io.q_buf.get(), {1, half, cfg->head_dim}, offset, "q dest"); - } - io.q_live = Window(io.q_buf.get(), {1, cfg->q_dim()}, 0, "q live"); - io.k_live = Window(io.k_buf.get(), {1, cfg->kv_dim()}, 0, "k live"); - io.mha_out = Tensor({1, cfg->q_dim()}, "mha"); - io.gated = Tensor({1, cfg->q_dim()}, "gated"); - io.k_cache = Tensor({cfg->kv_heads, kMaxSequenceLength, cfg->head_dim}, "k cache"); - io.v_cache = Tensor({cfg->kv_heads, kMaxSequenceLength, cfg->head_dim}, "v cache"); - io.v_rows.resize(static_cast(kMaxSequenceLength)); - } - void* cos_view = nullptr; - void* sin_view = nullptr; - const std::int64_t rope_shape[] = {kMaxSequenceLength, cfg->rope_dim / 2}; - api->Check(api->functions().create_host_view(ryzenai_corelib_data_type_fp32, rope_shape, 2, - cos_table.data(), &cos_view), - "ryzenai_corelib_create_host_view cos"); - api->Check(api->functions().create_host_view(ryzenai_corelib_data_type_fp32, rope_shape, 2, - sin_table.data(), &sin_view), - "ryzenai_corelib_create_host_view sin"); - cos = UniqueHostView(api, cos_view); - sin = UniqueHostView(api, sin_view); - } - - void BuildRope() { - const int half = cfg->rope_dim / 2; - cos_table.resize(static_cast(kMaxSequenceLength) * half); - sin_table.resize(cos_table.size()); - for (int pos = 0; pos < kMaxSequenceLength; ++pos) { - for (int i = 0; i < half; ++i) { - const double angle = pos / std::pow(cfg->rope_theta, (2.0 * i) / cfg->rope_dim); - const auto at = static_cast(pos) * half + i; - cos_table[at] = static_cast(std::cos(angle)); - sin_table[at] = static_cast(std::sin(angle)); - } - } - } - - void Reset() { - position = 0; - const auto cache_count = static_cast(cfg->kv_heads) * kMaxSequenceLength * cfg->head_dim; - for (int layer = 0; layer < cfg->layers; ++layer) { - if (cfg->is_full(layer)) { - Zero(full_io[layer].k_cache.get(), cache_count, "zero k cache"); - Zero(full_io[layer].v_cache.get(), cache_count, "zero v cache"); - Zero(full_io[layer].q_buf.get(), - static_cast(q_vectors) * cfg->head_dim, "zero q"); - Zero(full_io[layer].k_buf.get(), - static_cast(k_vectors) * cfg->head_dim, "zero k"); - } else { - auto& io = linear_io[layer]; - Zero(io.qkv.get(), static_cast(input_rows) * cfg->conv_channels(), "zero qkv"); - const auto conv_count = static_cast(input_rows) * cfg->conv_channels(); - const auto state_count = static_cast(cfg->lin_v_heads) * cfg->lin_v_dim * cfg->lin_k_dim; - for (int side = 0; side < 2; ++side) { - Zero(io.conv[side].get(), conv_count, "zero conv"); - Zero(io.state[side].get(), state_count, "zero state"); - } - parity[static_cast(layer)] = 0; - } - } - } - - void SubmitLinear(int layer, void* hidden) { - auto& weights = linear[layer]; - auto& io = linear_io[layer]; - const auto side = parity[static_cast(layer)]; - parity[static_cast(layer)] = 1 - side; - const auto& f = api->functions(); - api->Check(f.matmul(stream.get(), hidden, weights.qkv.get(), io.qkv_row.get()), "qkv"); - api->Check(f.matmul_act(stream.get(), ryzenai_corelib_matmul_activation_sigmoid, hidden, - weights.beta.get(), nullptr, io.beta.get()), - "beta"); - api->Check(f.matmul_act(stream.get(), ryzenai_corelib_matmul_activation_softplus_mulbcast, hidden, - weights.alpha.get(), weights.neg_a.get(), io.g.get()), - "alpha"); - api->Check(f.dwconv(stream.get(), io.qkv.get(), io.conv[side].get(), 1, weights.conv.get(), - io.dw_q.get(), io.dw_k.get(), io.dw_v.get(), io.conv[1 - side].get()), - "dwconv"); - api->Check(f.linear_attention(stream.get(), &lin_desc, io.state[side].get(), io.dw_q.get(), io.dw_k.get(), - io.dw_v.get(), io.g_live.get(), io.beta_live.get(), 1, io.lin_out.get(), - io.state[1 - side].get()), - "linear attention"); - api->Check(f.rmsnorm(stream.get(), io.lin_out.get(), weights.norm.get(), io.lin_out.get()), - "ssm norm"); - api->Check(f.matmul_act(stream.get(), ryzenai_corelib_matmul_activation_silu_mul, hidden, - weights.z.get(), io.lin_flat.get(), io.z_out.get()), - "z"); - api->Check(f.matmul(stream.get(), io.z_out.get(), weights.out.get(), attn[layer].get()), - "linear out"); - } - - void SubmitFull(int layer, void* hidden) { - auto& weights = full[layer]; - auto& io = full_io[layer]; - const auto& f = api->functions(); - auto& v_row = io.v_rows[static_cast(position)]; - if (!v_row) { - v_row = Window(io.v_cache.get(), {cfg->kv_heads, 1, cfg->head_dim}, - static_cast(position) * cfg->head_dim, "v row", - {kMaxSequenceLength * cfg->head_dim, cfg->head_dim, 1}); - } - api->Check(f.matmul(stream.get(), hidden, weights.q.get(), io.q_raw.get()), "q"); - api->Check(f.matmul(stream.get(), hidden, weights.k.get(), io.k_live.get()), "k"); - api->Check(f.matmul(stream.get(), hidden, weights.v.get(), v_row.get()), "v"); - for (int part = 0; part < 2; ++part) { - api->Check(f.rmsnorm(stream.get(), io.q_src[part].get(), weights.q_norm.get(), - io.q_dst[part].get()), - "q norm"); - } - api->Check(f.rmsnorm(stream.get(), io.k_buf.get(), weights.k_norm.get(), io.k_buf.get()), "k norm"); - api->Check(f.flat_mha(stream.get(), &mha_desc, io.q_live.get(), io.k_live.get(), position, cos.get(), - sin.get(), io.k_cache.get(), io.v_cache.get(), io.mha_out.get()), - "flat mha"); - api->Check(f.matmul_act(stream.get(), ryzenai_corelib_matmul_activation_sigmoid_mul, hidden, - weights.gate.get(), io.mha_out.get(), io.gated.get()), - "output gate"); - api->Check(f.matmul(stream.get(), io.gated.get(), weights.o.get(), attn[layer].get()), "o proj"); - } - - buffer StepBody(int token) { - if (position >= kMaxSequenceLength) - throw std::out_of_range("Qwen3.5 position is past the KV cache"); - if (static_cast(position) >= max_length) - throw std::out_of_range("Qwen3.5 request exceeds configured context capacity"); - const int id = token; - const auto embedded = rai::DecodeQ8Rows(package->Embedding(), std::span(&id, 1)); - bool submitted = false; - try { - Write(res_in.get(), embedded.data(), embedded.size(), "embedding"); - api->Check(api->functions().rmsnorm(stream.get(), res_in.get(), norm0.get(), h_in.get()), - "input norm"); - submitted = true; - void* residual = res_in.get(); - void* hidden = h_in.get(); - for (int layer = 0; layer < cfg->layers; ++layer) { - if (cfg->is_full(layer)) SubmitFull(layer, hidden); - else SubmitLinear(layer, hidden); - const auto& f = api->functions(); - void* mlp = cfg->is_full(layer) ? full[layer].mlp.get() : linear[layer].mlp.get(); - api->Check(f.ssmlp(stream.get(), attn[layer].get(), residual, mlp, res[layer].get(), - h[layer].get()), - "ssmlp"); - residual = res[layer].get(); - hidden = h[layer].get(); - } - api->Check(api->functions().matmul(stream.get(), hidden, lm_head.get(), logits.get()), "lm_head"); - api->Check(api->functions().stream_synchronize(stream.get()), - "ryzenai_corelib_stream_synchronize"); - ++position; - buffer out(static_cast(cfg->vocab)); - api->Check(api->functions().tensor_read(logits.get(), ryzenai_corelib_data_type_bf16, out.data(), - out.size(), 0), - "ryzenai_corelib_tensor_read logits"); - return out; - } catch (...) { - if (submitted) { - (void)api->functions().stream_synchronize(stream.get()); - poisoned = true; - position = 0; - saved.reset(); - history.clear(); - } - throw; - } - } - - std::shared_ptr package; - std::shared_ptr runtime; - std::shared_ptr api; - const Qwen35Config* cfg = nullptr; - float cfg_eps = 0; - std::vector epsilon_bf16; - ryzenai_corelib_linear_attention_bf16_desc lin_desc{}; - ryzenai_corelib_flat_mha_bf16_desc mha_desc{}; - int beta_width = 0; - int alpha_width = 0; - int q_vectors = 0; - int k_vectors = 0; - int input_rows = 0; - std::vector cos_table, sin_table, zeros; - std::vector history; - std::vector parity; - - UniqueStream stream; - UniqueRmsNormWeights norm0; - UniqueMatMulWeights lm_head; - UniqueHostView cos, sin; - std::vector full; - std::vector linear; - UniqueTensor res_in, h_in, logits; - std::vector res, h, attn; - std::vector full_io; - std::vector linear_io; -}; - -qwen35_rai::qwen35_rai(LM_Config, std::shared_ptr package, - std::shared_ptr runtime, std::uint32_t max_length) - : impl_(std::make_unique(std::move(package), std::move(runtime), max_length)) {} -qwen35_rai::~qwen35_rai() = default; - -buffer qwen35_rai::forward(int id) { return impl_->Forward(id); } -buffer qwen35_rai::prefill(std::vector& ids, void*) { return impl_->Prefill(ids); } - -void qwen35_rai::set_context_length(int n) { impl_->Rewind(n); } - -void qwen35_rai::load_weights(Q4NX&) { - impl_->usable(); - throw std::runtime_error("Qwen3.5 rai weights are loaded only from GGUF"); -} - -void qwen35_rai::update_max_length(std::uint32_t n) { - impl_->usable(); - if (!n || n > kMaxSequenceLength || n < static_cast(impl_->position)) - throw std::out_of_range("Qwen3.5 maximum length is invalid"); - impl_->max_length = n; -} - -void qwen35_rai::clear_context() { impl_->Clear(); } - -buffer qwen35_rai::get_k_cache(int layer, int index) { return impl_->ReadCache(true, layer, index); } -buffer qwen35_rai::get_v_cache(int layer, int index) { return impl_->ReadCache(false, layer, index); } - -int qwen35_rai::get_current_context_length() { - impl_->usable(); - return impl_->position; -} - -int qwen35_rai::checkpoint() { - impl_->usable(); - impl_->saved = impl_->position; - return impl_->position; -} - -int qwen35_rai::restore() { - impl_->usable(); - if (!impl_->saved) return -1; - impl_->Rewind(*impl_->saved); - return impl_->position; -} - -bool qwen35_rai::poisoned() const noexcept { return impl_ && impl_->poisoned; } - -} // namespace flm::qwen35 diff --git a/src/common/models/qwen35/rai/aie_next/qwen35_rai_backend.cpp b/src/common/models/qwen35/rai/aie_next/qwen35_rai_backend.cpp index cc0083453..5671fc357 100644 --- a/src/common/models/qwen35/rai/aie_next/qwen35_rai_backend.cpp +++ b/src/common/models/qwen35/rai/aie_next/qwen35_rai_backend.cpp @@ -1,79 +1,20 @@ +/// \file qwen35_rai_backend.cpp +/// \brief Dense Qwen3.5 on corelib's "qwen35" model #include "models/qwen35/rai/aie_next/qwen35_rai_backend.hpp" -#include "models/qwen35/rai/aie_next/qwen35_rai.hpp" -#include "rai/corelib_runtime.hpp" -#include "utils/utils.hpp" - -#include -#include -#include -#include +#include "rai/corelib_model_backend.hpp" namespace flm::qwen35 { -namespace { - -std::string GgufName(const nlohmann::ordered_json& model_info) { - std::string found; - if (model_info.contains("files") && model_info["files"].is_array()) { - for (const auto& file : model_info["files"]) { - if (!file.is_string()) continue; - const auto name = file.get(); - if (name.size() > 5 && name.ends_with(".gguf")) { - if (!found.empty()) - throw std::runtime_error("Qwen3.5 rai entry lists more than one .gguf file"); - found = name; - } - } - } - if (found.empty()) throw std::runtime_error("Qwen3.5 rai entry lists no .gguf file"); - return found; -} - -class RaiBackend final : public flm::backend::ModelBackend { -public: - explicit RaiBackend(const flm::backend::BackendContext& context) { - if (!context.config) throw std::runtime_error("Qwen3.5 rai backend needs an LM_Config"); - if (context.enable_preemption) - throw std::invalid_argument("Qwen3.5 rai does not support preemption"); - if (context.context_length < 1 || context.context_length > kRaiContextLimit) - throw std::out_of_range("Qwen3.5 rai context length must be in 1..4096"); - - auto package = Qwen35Gguf::Open( - std::filesystem::path(context.model_path) / GgufName(context.model_info)); - runtime_ = corelib::CorelibRuntime::GetOrCreate(utils::get_executable_directory()); - auto engine = std::make_unique(*context.config, std::move(package), runtime_, - context.context_length); - engine->clear_context(); - engine_ = std::move(engine); - } - - ~RaiBackend() override { engine_.reset(); } - - causal_lm& engine() override { return *engine_; } - std::string id() const override { return flm::backend::kRaiBackendId; } - std::string detail() const override { - return runtime_ ? runtime_->loaded_library_path().string() : std::string(); - } - std::uint32_t max_decode_length() const override { return kRaiDecodeLimit; } - bool supports_preemption() const override { return false; } - bool forwards_past_eos() const override { return false; } - bool poisoned() const noexcept override { return engine_ && engine_->poisoned(); } - std::optional> forced_eos_ids() const override { - return std::vector(kEosIds.begin(), kEosIds.end()); - } - -private: - std::shared_ptr runtime_; - std::unique_ptr engine_; -}; - -} // namespace +/// \note The linear-attention state can neither rewind nor skip and corelib +/// keeps it outside any state slot, so going back means replaying. flm::backend::BackendFactory rai_factory() { - return [](const flm::backend::BackendContext& context) - -> std::unique_ptr { - return std::make_unique(context); - }; + return rai::Factory({ + .model = {"qwen35", "Qwen3.5", kRaiContextLimit, rai::Continuation::kAnywhere, + rai::Rewind::kReplay}, + .context_rule = "1..4096", + .eos_ids = {248046, 248044}, + }); } } // namespace flm::qwen35 diff --git a/src/common/models/qwen35/rai/aie_next/qwen35_rai_gguf.cpp b/src/common/models/qwen35/rai/aie_next/qwen35_rai_gguf.cpp deleted file mode 100644 index faf84b68b..000000000 --- a/src/common/models/qwen35/rai/aie_next/qwen35_rai_gguf.cpp +++ /dev/null @@ -1,355 +0,0 @@ -#include "models/qwen35/rai/aie_next/qwen35_rai_gguf.hpp" - -#include -#include -#include -#include -#include -#include - -namespace flm::qwen35 { -namespace { - -constexpr int kBlockWeights = 32; -constexpr int kBlockBytes = 34; -constexpr float kSteps = 15.0f; - -std::uint64_t NeedUnsigned(const rai::GgufFile& file, std::string_view key) { - if (!file.HasMetadata(key)) - throw std::runtime_error(std::string("the GGUF has no ") + std::string(key)); - return file.Unsigned(key); -} - -double NeedNumber(const rai::GgufFile& file, std::string_view key) { - if (!file.HasMetadata(key)) - throw std::runtime_error(std::string("the GGUF has no ") + std::string(key)); - return file.Number(key); -} - -std::vector Expand(const std::vector& order, int width) { - std::vector out(order.size() * static_cast(width)); - for (std::size_t head = 0; head < order.size(); ++head) { - for (int i = 0; i < width; ++i) - out[head * static_cast(width) + static_cast(i)] = - order[head] * width + i; - } - return out; -} - -float HalfToFloat(std::uint16_t half) { - const std::uint32_t sign = static_cast(half & 0x8000u) << 16; - const std::uint32_t exponent = (half >> 10) & 0x1fu; - std::uint32_t fraction = half & 0x03ffu; - std::uint32_t bits = 0; - if (exponent == 0) { - if (fraction == 0) { - bits = sign; - } else { - int shift = 0; - while ((fraction & 0x0400u) == 0) { - fraction <<= 1; - ++shift; - } - fraction &= 0x03ffu; - bits = sign | (static_cast(127 - 14 - shift) << 23) | (fraction << 13); - } - } else if (exponent == 0x1fu) { - bits = sign | 0x7f800000u | (fraction << 13); - } else { - bits = sign | ((exponent + (127 - 15)) << 23) | (fraction << 13); - } - return std::bit_cast(bits); -} - -float Bf16Up(float value) { - const auto bits = std::bit_cast(value); - auto truncated = static_cast(bits >> 16); - if ((bits & 0xFFFFu) != 0) truncated = static_cast(truncated + 1); - return std::bit_cast(static_cast(truncated) << 16); -} - -std::uint16_t FloatToHalf(float value) { - const auto bits = std::bit_cast(value); - const std::uint32_t sign = (bits >> 16) & 0x8000u; - auto exponent = static_cast((bits >> 23) & 0xffu) - 127 + 15; - std::uint32_t mantissa = bits & 0x7fffffu; - if (exponent <= 0) { - if (exponent < -10) return static_cast(sign); - mantissa |= 0x800000u; - const auto shift = static_cast(14 - exponent); - std::uint32_t half = mantissa >> shift; - const auto remainder = mantissa & ((1u << shift) - 1u); - if (remainder > (1u << (shift - 1)) || - (remainder == (1u << (shift - 1)) && (half & 1u))) - ++half; - return static_cast(sign | half); - } - if (exponent >= 31) return static_cast(sign | 0x7c00u); - std::uint32_t half = mantissa >> 13; - const auto remainder = mantissa & 0x1fffu; - if (remainder > 0x1000u || (remainder == 0x1000u && (half & 1u))) { - ++half; - if (half == 0x400u) { - half = 0; - ++exponent; - if (exponent >= 31) return static_cast(sign | 0x7c00u); - } - } - return static_cast(sign | (static_cast(exponent) << 10) | half); -} - -std::vector Dequantize(std::span blocks, int rows, int cols) { - const int blocks_per_row = cols / kBlockWeights; - if (static_cast(blocks.size()) != rows * blocks_per_row * kBlockBytes) - throw std::runtime_error("Q8_0 byte length does not match the matrix being refit"); - std::vector out(static_cast(rows) * static_cast(cols)); - for (int row = 0; row < rows; ++row) { - for (int block = 0; block < blocks_per_row; ++block) { - const auto* encoded = - blocks.data() + (static_cast(row) * blocks_per_row + block) * kBlockBytes; - std::uint16_t scale_bits = 0; - std::memcpy(&scale_bits, encoded, sizeof(scale_bits)); - const float scale = HalfToFloat(scale_bits); - for (int element = 0; element < kBlockWeights; ++element) { - const auto code = static_cast( - std::to_integer(encoded[2 + element])); - out[static_cast(row) * cols + block * kBlockWeights + element] = - scale * static_cast(code); - } - } - } - return out; -} - -} // namespace - -Qwen35Gguf::Qwen35Gguf(std::shared_ptr file, std::filesystem::path path, - Qwen35Config config, std::string lm_head, rai::GgufTensor embedding, - std::vector order) - : file_(std::move(file)), path_(std::move(path)), config_(config), lm_head_(std::move(lm_head)), - embedding_(std::move(embedding)), order_(std::move(order)) {} - -std::shared_ptr Qwen35Gguf::Open(const std::filesystem::path& path) { - auto file = rai::GgufFile::Open(path, [](std::string_view key) { - return key == "general.architecture" || key.starts_with("qwen35."); - }); - const auto arch = file->String("general.architecture"); - if (arch != "qwen35") - throw std::runtime_error("general.architecture is " + arch + - "; this engine reads the dense Qwen3.5 conversion"); - - Qwen35Config config; - config.layers = static_cast(NeedUnsigned(*file, "qwen35.block_count")); - config.hidden = static_cast(NeedUnsigned(*file, "qwen35.embedding_length")); - config.intermediate = static_cast(NeedUnsigned(*file, "qwen35.feed_forward_length")); - config.q_heads = static_cast(NeedUnsigned(*file, "qwen35.attention.head_count")); - config.kv_heads = static_cast(NeedUnsigned(*file, "qwen35.attention.head_count_kv")); - config.head_dim = static_cast(NeedUnsigned(*file, "qwen35.attention.key_length")); - const auto value_dim = static_cast(NeedUnsigned(*file, "qwen35.attention.value_length")); - if (value_dim != config.head_dim) - throw std::runtime_error("key_length and value_length differ"); - config.rope_dim = static_cast(NeedUnsigned(*file, "qwen35.rope.dimension_count")); - config.rope_theta = NeedNumber(*file, "qwen35.rope.freq_base"); - config.eps = NeedNumber(*file, "qwen35.attention.layer_norm_rms_epsilon"); - config.full_every = static_cast(NeedUnsigned(*file, "qwen35.full_attention_interval")); - config.conv_width = static_cast(NeedUnsigned(*file, "qwen35.ssm.conv_kernel")); - config.lin_k_dim = static_cast(NeedUnsigned(*file, "qwen35.ssm.state_size")); - config.lin_k_heads = static_cast(NeedUnsigned(*file, "qwen35.ssm.group_count")); - config.lin_v_heads = static_cast(NeedUnsigned(*file, "qwen35.ssm.time_step_rank")); - const auto inner = static_cast(NeedUnsigned(*file, "qwen35.ssm.inner_size")); - if (config.full_every <= 0 || config.lin_v_heads <= 0 || config.lin_k_heads <= 0 || - config.lin_v_heads % config.lin_k_heads != 0 || inner % config.lin_v_heads != 0) - throw std::runtime_error("Qwen3.5 linear-attention metadata does not divide evenly"); - config.lin_v_dim = inner / config.lin_v_heads; - - const auto table = file->Tensor("token_embd.weight"); - if (table.ggml_type != rai::kGgmlTypeQ8_0 || table.shape.size() != 2 || - table.shape[1] != config.hidden) - throw std::runtime_error("token_embd.weight is not Q8_0 [vocab, hidden]"); - config.vocab = static_cast(table.shape[0]); - config.tied_lm_head = !file->HasTensor("output.weight"); - const std::string lm_head = config.tied_lm_head ? "token_embd" : "output"; - if (file->Tensor(lm_head + ".weight").ggml_type != rai::kGgmlTypeQ8_0 || - file->Tensor("blk.0.attn_qkv.weight").ggml_type != rai::kGgmlTypeQ8_0) - throw std::runtime_error("this engine packs Q8_0, which corelib refits to group 64"); - - const int per = config.lin_v_heads / config.lin_k_heads; - std::vector order(static_cast(config.lin_v_heads)); - for (int head = 0; head < config.lin_v_heads; ++head) - order[static_cast(head)] = (head % per) * config.lin_k_heads + head / per; - - return std::shared_ptr(new Qwen35Gguf( - std::move(file), path, config, lm_head, table, std::move(order))); -} - -std::span Qwen35Gguf::Blocks(std::string_view base, std::int64_t k, - std::int64_t n) const { - const std::int64_t shape[] = {n, k}; - const auto tensor = file_->RequireQ8(std::string(base) + ".weight", shape); - return tensor.bytes; -} - -std::span Qwen35Gguf::F32(std::string_view name, - std::span shape) const { - return file_->RequireF32(name, shape).values; -} - -std::vector Qwen35Gguf::TakeRows(std::span raw, std::int64_t n, - std::int64_t k, - std::span rows) const { - if (k % kBlockWeights != 0) throw std::runtime_error("Q8_0 width is not a whole number of blocks"); - const auto row_bytes = static_cast(k / kBlockWeights * kBlockBytes); - if (raw.size() != static_cast(n) * row_bytes) - throw std::runtime_error("Q8_0 row table is the wrong length"); - std::vector out(rows.size() * row_bytes); - for (std::size_t i = 0; i < rows.size(); ++i) { - const auto row = rows[i]; - if (row < 0 || row >= n) throw std::runtime_error("Q8_0 row reorder is out of range"); - std::memcpy(out.data() + i * row_bytes, raw.data() + static_cast(row) * row_bytes, - row_bytes); - } - return out; -} - -std::vector Qwen35Gguf::ChannelOrder() const { - const auto qk = config_.lin_qk_dim() * 2; - auto order = Expand(order_, config_.lin_v_dim); - std::vector rows(static_cast(config_.conv_channels())); - for (int i = 0; i < qk; ++i) rows[static_cast(i)] = i; - for (std::size_t i = 0; i < order.size(); ++i) - rows[static_cast(qk) + i] = qk + order[i]; - return rows; -} - -std::vector Qwen35Gguf::QkvBlocks(int layer) const { - const auto base = "blk." + std::to_string(layer) + ".attn_qkv"; - const auto raw = Blocks(base, config_.hidden, config_.conv_channels()); - const auto rows = ChannelOrder(); - return TakeRows(raw, config_.conv_channels(), config_.hidden, rows); -} - -std::vector Qwen35Gguf::ZBlocks(int layer) const { - const auto raw = Blocks("blk." + std::to_string(layer) + ".attn_gate", config_.hidden, - config_.lin_value_dim()); - const auto rows = Expand(order_, config_.lin_v_dim); - return TakeRows(raw, config_.lin_value_dim(), config_.hidden, rows); -} - -std::vector Qwen35Gguf::GateBlocks(int layer, std::string_view which) const { - const auto raw = Blocks("blk." + std::to_string(layer) + ".ssm_" + std::string(which), - config_.hidden, config_.lin_v_heads); - return TakeRows(raw, config_.lin_v_heads, config_.hidden, order_); -} - -std::vector Qwen35Gguf::OutBlocks(int layer) const { - const auto raw = Blocks("blk." + std::to_string(layer) + ".ssm_out", config_.lin_value_dim(), - config_.hidden); - const auto per = config_.lin_v_dim / kBlockWeights; - const auto head_bytes = static_cast(per * kBlockBytes); - const auto row_bytes = head_bytes * static_cast(config_.lin_v_heads); - if (raw.size() != static_cast(config_.hidden) * row_bytes) - throw std::runtime_error("ssm_out Q8_0 length does not match its heads"); - std::vector out(raw.size()); - for (int row = 0; row < config_.hidden; ++row) { - const auto* src = raw.data() + static_cast(row) * row_bytes; - auto* dst = out.data() + static_cast(row) * row_bytes; - for (int head = 0; head < config_.lin_v_heads; ++head) { - std::memcpy(dst + static_cast(head) * head_bytes, - src + static_cast(order_[static_cast(head)]) * head_bytes, - head_bytes); - } - } - return out; -} - -std::vector Qwen35Gguf::PerHead(int layer, std::string_view name) const { - const std::int64_t shape[] = {config_.lin_v_heads}; - const auto values = F32("blk." + std::to_string(layer) + "." + std::string(name), shape); - std::vector out(order_.size()); - for (std::size_t head = 0; head < order_.size(); ++head) - out[head] = values[static_cast(order_[head])]; - return out; -} - -std::vector Qwen35Gguf::ConvTaps(int layer) const { - const std::int64_t shape[] = {config_.conv_channels(), config_.conv_width}; - const auto values = F32("blk." + std::to_string(layer) + ".ssm_conv1d.weight", shape); - const auto rows = ChannelOrder(); - std::vector out(values.size()); - for (std::size_t row = 0; row < rows.size(); ++row) { - const auto src = static_cast(rows[row]) * static_cast(config_.conv_width); - const auto dst = row * static_cast(config_.conv_width); - std::copy_n(values.begin() + static_cast(src), config_.conv_width, - out.begin() + static_cast(dst)); - } - return out; -} - -OnnxQ8 Q8ToOnnx(std::span blocks, int k, int n, int group) { - if (k % group != 0) throw std::runtime_error("K is not a whole number of quantization groups"); - auto weights = Dequantize(blocks, n, k); - std::vector transposed(static_cast(k) * static_cast(n)); - for (int row = 0; row < n; ++row) { - for (int col = 0; col < k; ++col) - transposed[static_cast(col) * n + row] = - weights[static_cast(row) * k + col]; - } - - const int groups = k / group; - std::vector scales(static_cast(groups) * static_cast(n)); - std::vector zeros(scales.size()); - std::vector values(static_cast(k) * static_cast(n)); - for (int group_index = 0; group_index < groups; ++group_index) { - for (int col = 0; col < n; ++col) { - float lo = 0.0f; - float hi = 0.0f; - for (int row = 0; row < group; ++row) { - const float value = - transposed[(static_cast(group_index) * group + row) * n + col]; - lo = std::min(lo, value); - hi = std::max(hi, value); - } - float scale = Bf16Up((hi - lo) / kSteps); - std::uint8_t zero = 8; - if (scale == 0.0f) { - scale = 1.0f; - } else { - zero = static_cast( - std::clamp(std::floor(-lo / scale + 0.5f), 0.0f, kSteps)); - } - scales[static_cast(group_index) * n + col] = scale; - zeros[static_cast(group_index) * n + col] = zero; - for (int row = 0; row < group; ++row) { - const auto at = (static_cast(group_index) * group + row) * n + col; - const float code = std::floor(transposed[at] / scale + 0.5f) + static_cast(zero); - values[at] = static_cast(std::clamp(code, 0.0f, kSteps)); - } - } - } - - OnnxQ8 onnx; - onnx.qweight.resize(static_cast(n) * static_cast(k / 2)); - onnx.scales.resize(static_cast(n) * static_cast(groups)); - onnx.qzeros.resize(static_cast(n) * static_cast(groups / 2)); - for (int row = 0; row < n; ++row) { - for (int col = 0; col < k; col += 2) { - const auto low = values[static_cast(col) * n + row]; - const auto high = values[static_cast(col + 1) * n + row]; - onnx.qweight[static_cast(row) * (k / 2) + col / 2] = - static_cast(low | (high << 4)); - } - for (int group_index = 0; group_index < groups; ++group_index) { - onnx.scales[static_cast(row) * groups + group_index] = - FloatToHalf(scales[static_cast(group_index) * n + row]); - } - for (int group_index = 0; group_index < groups; group_index += 2) { - const auto low = zeros[static_cast(group_index) * n + row]; - const auto high = zeros[static_cast(group_index + 1) * n + row]; - onnx.qzeros[static_cast(row) * (groups / 2) + group_index / 2] = - static_cast(low | (high << 4)); - } - } - return onnx; -} - -} // namespace flm::qwen35 diff --git a/src/common/models/qwen36/rai/aie_next/qwen36_rai.cpp b/src/common/models/qwen36/rai/aie_next/qwen36_rai.cpp deleted file mode 100644 index caef9bc8a..000000000 --- a/src/common/models/qwen36/rai/aie_next/qwen36_rai.cpp +++ /dev/null @@ -1,289 +0,0 @@ -#include "models/qwen36/rai/aie_next/qwen36_rai.hpp" - -#include "rai/corelib_object.hpp" -#include "rai/host_ops.hpp" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace flm::qwen36 { -namespace { - -using flm::corelib::CorelibApi; -using flm::corelib::CorelibRuntime; -using flm::corelib::UniqueModel; -using flm::corelib::UniqueModelSnapshot; -using flm::corelib::UniqueModelWeights; - -} // namespace - -class qwen36_rai::Impl { -public: - Impl(std::filesystem::path gguf, std::shared_ptr runtime, - std::uint32_t maximum) - : runtime(std::move(runtime)), - api(this->runtime ? this->runtime->api() : nullptr), - max_length(maximum) { - if (gguf.empty()) throw std::invalid_argument("Qwen3.6 GGUF path is empty"); - if (!this->runtime || !api) throw std::invalid_argument("corelib runtime is null"); - if (!maximum || maximum > kContextLimit) - throw std::invalid_argument("Qwen3.6 maximum length must be in 1.." + - std::to_string(kContextLimit)); - - std::cout << "[FLM] Loading Qwen3.6-35B-A3B from Q8_0 GGUF" << std::endl; - auto lease = this->runtime->AcquireExecution(); - const auto path = gguf.string(); - void* raw_weights = nullptr; - api->Check(api->functions().model_weights_create_path("qwen36", path.c_str(), &raw_weights), - "ryzenai_corelib_model_weights_create_path"); - weights = UniqueModelWeights(api, raw_weights); - Recreate(maximum); - logits.resize(vocab); - } - - buffer Forward(int id) { - usable(); - auto lease = runtime->AcquireExecution(); - Run(&id, 1, true); - history.push_back(id); - return Logits(); - } - - buffer Prefill(const std::vector& ids, void* payload) { - usable(); - if (payload) - throw std::invalid_argument("Qwen3.6 rai does not run a vision payload"); - if (ids.empty()) throw std::invalid_argument("Qwen3.6 request contains no token IDs"); - auto lease = runtime->AcquireExecution(); - Run(ids.data(), ids.size(), true); - history.insert(history.end(), ids.begin(), ids.end()); - return Logits(); - } - - void Clear() { - usable(); - auto lease = runtime->AcquireExecution(); - ResetModel(); - history.clear(); - saved.reset(); - } - - void Rewind(int to) { - usable(); - if (to < 0 || static_cast(to) > max_length) - throw std::out_of_range("Qwen3.6 context length is out of range"); - if (static_cast(to) > history.size()) - throw std::runtime_error("Qwen3.6 cannot rewind past the tokens it has run"); - if (to == position) return; - auto lease = runtime->AcquireExecution(); - if (saved && saved->position == to && history.size() >= saved->history.size() && - std::equal(saved->history.begin(), saved->history.end(), history.begin())) { - Load(saved->snapshot, to); - history = saved->history; - return; - } - history.resize(static_cast(to)); - ResetModel(); - if (to > 0) Run(history.data(), history.size(), false); - } - - int Checkpoint() { - usable(); - if (Matches(saved)) return position; - auto lease = runtime->AcquireExecution(); - saved = Saved{position, history, Save()}; - return position; - } - - int Restore() { - usable(); - if (!saved) return -1; - if (Matches(saved)) return position; - auto lease = runtime->AcquireExecution(); - Load(saved->snapshot, saved->position); - history = saved->history; - return position; - } - - void Grow(std::uint32_t n) { - if (n <= capacity) return; - auto lease = runtime->AcquireExecution(); - UniqueModelSnapshot current; - if (position > 0) current = Save(); - Recreate(n); - if (current) Load(current, position); - } - - void usable() const { - if (poisoned) throw std::runtime_error("Qwen3.6 corelib engine is poisoned"); - } - - struct Saved { - int position = 0; - std::vector history; - UniqueModelSnapshot snapshot; - }; - - bool poisoned = false; - - std::shared_ptr runtime; - std::shared_ptr api; - UniqueModelWeights weights; - UniqueModel model; - std::uint32_t capacity = 0; - std::uint32_t max_length = 0; - std::size_t vocab = 0; - int position = 0; - std::vector history; - std::optional saved; - std::vector logits; - -private: - /// \brief a model over `weights` holding at least `n` positions - /// \note The capacity is fixed for a model's life; `weights` outlive it, so - /// a larger one is a new model over the same packed weights. - void Recreate(std::uint32_t n) { - ryzenai_corelib_model_create_options options{}; - options.size = sizeof(options); - options.max_sequence_length = static_cast(n); - void* raw = nullptr; - api->Check(api->functions().model_create("qwen36", weights.get(), &options, &raw), - "ryzenai_corelib_model_create"); - UniqueModel created(api, raw); - ryzenai_corelib_model_info info{}; - api->Check(api->functions().model_info_get(created.get(), &info), - "ryzenai_corelib_model_info_get"); - if (info.vocab_size <= 0 || info.context_length < static_cast(n)) { - throw std::runtime_error( - "Qwen3.6 model_info is vocab " + std::to_string(info.vocab_size) + - " context " + std::to_string(info.context_length) + - "; this engine needs vocab > 0 and context >= " + std::to_string(n)); - } - model = std::move(created); - capacity = static_cast(info.context_length); - vocab = static_cast(info.vocab_size); - } - - void Run(const int* ids, std::size_t rows, bool want_logits) { - if (static_cast(position) + rows > max_length) - throw std::out_of_range("Qwen3.6 request exceeds configured context capacity"); - const std::vector tokens(ids, ids + rows); - ryzenai_corelib_model_forward_options options{}; - options.position = RYZENAI_CORELIB_POSITION_CONTINUE; - if (want_logits) { - options.logits = logits.data(); - options.logits_capacity = logits.size(); - } - try { - api->Check(api->functions().model_forward( - model.get(), tokens.data(), nullptr, rows, nullptr, &options), - "ryzenai_corelib_model_forward"); - } catch (...) { - poisoned = true; - position = 0; - history.clear(); - saved.reset(); - throw; - } - position += static_cast(rows); - } - - UniqueModelSnapshot Save() { - void* raw = nullptr; - api->Check(api->functions().model_state_save(model.get(), &raw), - "ryzenai_corelib_model_state_save"); - return UniqueModelSnapshot(api, raw); - } - - /// \note corelib leaves the model reset when a restore fails partway, so a - /// failure leaves nothing this engine's position or history describe. - void Load(const UniqueModelSnapshot& snapshot, int at) { - try { - api->Check(api->functions().model_state_restore(model.get(), snapshot.get()), - "ryzenai_corelib_model_state_restore"); - } catch (...) { - poisoned = true; - position = 0; - history.clear(); - saved.reset(); - throw; - } - position = at; - } - - bool Matches(const std::optional& snapshot) const { - return snapshot && snapshot->position == position && snapshot->history == history; - } - - buffer Logits() const { - buffer out(vocab); - rai::F32ToBf16(logits, std::span( - reinterpret_cast(out.data()), vocab)); - return out; - } - - void ResetModel() { - api->Check(api->functions().model_reset(model.get()), "ryzenai_corelib_model_reset"); - position = 0; - } -}; - -qwen36_rai::qwen36_rai(std::filesystem::path gguf, - std::shared_ptr runtime, - std::uint32_t max_length) - : impl_(std::make_unique(std::move(gguf), std::move(runtime), max_length)) {} - -qwen36_rai::~qwen36_rai() = default; - -buffer qwen36_rai::forward(int id) { return impl_->Forward(id); } - -buffer qwen36_rai::prefill(std::vector& ids, void* payload) { - return impl_->Prefill(ids, payload); -} - -void qwen36_rai::set_context_length(int n) { impl_->Rewind(n); } - -void qwen36_rai::load_weights(Q4NX&) { - impl_->usable(); - throw std::runtime_error("Qwen3.6 rai weights are loaded only from GGUF"); -} - -void qwen36_rai::update_max_length(std::uint32_t n) { - impl_->usable(); - if (!n || n > kContextLimit || n < static_cast(impl_->position)) - throw std::out_of_range("Qwen3.6 maximum length is invalid"); - impl_->Grow(n); - if (impl_->saved && n < static_cast(impl_->saved->position)) - impl_->saved.reset(); - impl_->max_length = n; -} - -void qwen36_rai::clear_context() { impl_->Clear(); } - -buffer qwen36_rai::get_k_cache(int, int) { - throw std::runtime_error("Qwen3.6 rai does not expose the KV cache"); -} - -buffer qwen36_rai::get_v_cache(int, int) { - throw std::runtime_error("Qwen3.6 rai does not expose the KV cache"); -} - -int qwen36_rai::get_current_context_length() { - impl_->usable(); - return impl_->position; -} - -int qwen36_rai::checkpoint() { return impl_->Checkpoint(); } - -int qwen36_rai::restore() { return impl_->Restore(); } - -bool qwen36_rai::poisoned() const noexcept { return impl_ && impl_->poisoned; } - -} // namespace flm::qwen36 diff --git a/src/common/models/qwen36/rai/aie_next/qwen36_rai_backend.cpp b/src/common/models/qwen36/rai/aie_next/qwen36_rai_backend.cpp index 23d1138f4..db8aa0d66 100644 --- a/src/common/models/qwen36/rai/aie_next/qwen36_rai_backend.cpp +++ b/src/common/models/qwen36/rai/aie_next/qwen36_rai_backend.cpp @@ -1,109 +1,19 @@ +/// \file qwen36_rai_backend.cpp +/// \brief Qwen3.6-35B-A3B text on corelib's "qwen36" model #include "models/qwen36/rai/aie_next/qwen36_rai_backend.hpp" -#include "models/qwen36/rai/aie_next/qwen36_rai.hpp" -#include "rai/corelib_runtime.hpp" -#include "utils/utils.hpp" +#include "rai/corelib_model_backend.hpp" -#include -#include -#include -#include -#include #include -#include namespace flm::qwen36 { -namespace { - -std::string GgufName(const nlohmann::ordered_json& model_info) { - std::string found; - if (model_info.contains("files") && model_info["files"].is_array()) { - for (const auto& file : model_info["files"]) { - if (!file.is_string()) continue; - const auto name = file.get(); - if (name.size() > 5 && name.ends_with(".gguf")) { - if (!found.empty()) - throw std::runtime_error("Qwen3.6 rai entry lists more than one .gguf file"); - found = name; - } - } - } - if (found.empty()) throw std::runtime_error("Qwen3.6 rai entry lists no .gguf file"); - return found; -} - -// The published tokenizer_config names eos_token as a string and does not -// carry eos_token_id. The id is the added token of that same string. -std::vector EosIds(const std::filesystem::path& model) { - std::ifstream config_file(model / "tokenizer_config.json"); - if (!config_file) throw std::runtime_error("Qwen3.6 tokenizer_config.json is missing"); - const auto config = nlohmann::json::parse(config_file); - if (!config.contains("eos_token") || !config["eos_token"].is_string()) - throw std::runtime_error("Qwen3.6 tokenizer_config.json has no eos_token string"); - const auto eos = config["eos_token"].get(); - - std::ifstream tokenizer_file(model / "tokenizer.json"); - if (!tokenizer_file) throw std::runtime_error("Qwen3.6 tokenizer.json is missing"); - const auto tokenizer = nlohmann::json::parse(tokenizer_file); - std::vector ids; - if (tokenizer.contains("added_tokens") && tokenizer["added_tokens"].is_array()) { - for (const auto& token : tokenizer["added_tokens"]) { - if (token.value("content", std::string()) == eos && token.contains("id") && - token["id"].is_number_integer()) - ids.push_back(token["id"].get()); - } - } - if (ids.empty()) - throw std::runtime_error("Qwen3.6 tokenizer.json has no id for eos token " + eos); - return ids; -} - -class RaiBackend final : public flm::backend::ModelBackend { -public: - explicit RaiBackend(const flm::backend::BackendContext& context) { - if (context.enable_preemption) - throw std::invalid_argument("Qwen3.6 rai does not support preemption"); - if (context.context_length < 1 || context.context_length > kRaiContextLimit) - throw std::out_of_range("Qwen3.6 rai context length must be in 1.." + - std::to_string(kRaiContextLimit)); - context_length_ = context.context_length; - eos_ids_ = EosIds(context.model_path); - - runtime_ = corelib::CorelibRuntime::GetOrCreate(utils::get_executable_directory()); - auto engine = std::make_unique( - std::filesystem::path(context.model_path) / GgufName(context.model_info), - runtime_, context.context_length); - engine->clear_context(); - engine_ = std::move(engine); - } - - ~RaiBackend() override { engine_.reset(); } - - causal_lm& engine() override { return *engine_; } - std::string id() const override { return flm::backend::kRaiBackendId; } - std::string detail() const override { - return runtime_ ? runtime_->loaded_library_path().string() : std::string(); - } - std::uint32_t max_decode_length() const override { return context_length_ - 1; } - bool supports_preemption() const override { return false; } - bool forwards_past_eos() const override { return false; } - bool poisoned() const noexcept override { return engine_ && engine_->poisoned(); } - std::optional> forced_eos_ids() const override { return eos_ids_; } - -private: - std::uint32_t context_length_{}; - std::vector eos_ids_; - std::shared_ptr runtime_; - std::unique_ptr engine_; -}; - -} // namespace flm::backend::BackendFactory rai_factory() { - return [](const flm::backend::BackendContext& context) - -> std::unique_ptr { - return std::make_unique(context); - }; + return rai::Factory({ + .model = {"qwen36", "Qwen3.6", kRaiContextLimit, rai::Continuation::kAnywhere, + rai::Rewind::kSnapshot}, + .context_rule = "1.." + std::to_string(kRaiContextLimit), + }); } } // namespace flm::qwen36 diff --git a/src/common/rai/corelib_model.cpp b/src/common/rai/corelib_model.cpp new file mode 100644 index 000000000..31bb8cbb1 --- /dev/null +++ b/src/common/rai/corelib_model.cpp @@ -0,0 +1,352 @@ +/// \file corelib_model.cpp +#include "rai/corelib_model.hpp" + +#include "rai/corelib_object.hpp" +#include "rai/host_ops.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace flm::rai { + +using corelib::CorelibApi; +using corelib::CorelibRuntime; +using corelib::UniqueModel; +using corelib::UniqueModelSnapshot; +using corelib::UniqueModelWeights; + +struct CorelibModel::Impl { + Impl(ModelSpec model_spec, const std::filesystem::path& file, + std::shared_ptr owner, std::uint32_t maximum) + : spec(std::move(model_spec)), + runtime(std::move(owner)), + api(runtime ? runtime->api() : nullptr), + max_length(maximum) { + if (spec.rewind != Rewind::kSeek && spec.continuation != Continuation::kAnywhere) + throw std::invalid_argument(spec.display_name + + ": a model that cannot seek has to take a pass anywhere"); + if (file.empty()) throw std::invalid_argument(spec.display_name + " weights path is empty"); + if (!runtime || !api) throw std::invalid_argument("corelib runtime is null"); + if (!maximum || maximum > spec.context_limit) + throw std::invalid_argument(spec.display_name + " maximum length must be in 1.." + + std::to_string(spec.context_limit)); + + std::cout << "[FLM] Loading " << spec.display_name << " from " + << file.filename().string() << std::endl; + auto lease = runtime->AcquireExecution(); + const auto path = file.string(); + void* raw = nullptr; + api->Check(api->functions().model_weights_create_path(spec.corelib_name.c_str(), + path.c_str(), &raw), + "ryzenai_corelib_model_weights_create_path"); + weights = UniqueModelWeights(api, raw); + Recreate(maximum); + } + + struct Saved { + std::vector history; + /// \brief empty unless Rewind::kSnapshot + UniqueModelSnapshot snapshot; + }; + + ModelSpec spec; + std::shared_ptr runtime; + std::shared_ptr api; + UniqueModelWeights weights; + UniqueModel model; + std::uint32_t capacity = 0; + std::uint32_t max_length = 0; + std::size_t vocab = 0; + /// \brief every token the model holds, in order; its size is the position + std::vector history; + std::optional saved; + std::vector logits; + bool poisoned = false; + + int Position() const { return static_cast(history.size()); } + + void Usable() const { + if (poisoned) throw std::runtime_error(spec.display_name + " corelib engine is poisoned"); + } + + buffer Prefill(const std::vector& ids, void* payload) { + Usable(); + if (payload) + throw std::invalid_argument(spec.display_name + " rai does not run a vision payload"); + if (ids.empty()) + throw std::invalid_argument(spec.display_name + " request contains no token IDs"); + CheckIds(ids); + auto lease = runtime->AcquireExecution(); + Advance(ids, true); + return Logits(); + } + + buffer Forward(int id) { + Usable(); + CheckIds(std::span(&id, 1)); + auto lease = runtime->AcquireExecution(); + Advance(std::span(&id, 1), true); + return Logits(); + } + + // Not named Rewind: inside Impl that would hide the enum of the same name. + void RewindTo(int to) { + Usable(); + if (to < 0 || static_cast(to) > max_length) + throw std::out_of_range(spec.display_name + " context length is out of range"); + if (to > Position()) + throw std::runtime_error(spec.display_name + " cannot rewind past the tokens it has run"); + if (to == Position()) return; + auto lease = runtime->AcquireExecution(); + GoTo(std::vector(history.begin(), history.begin() + to)); + } + + int Checkpoint() { + Usable(); + if (saved && saved->history == history) return Position(); + auto lease = runtime->AcquireExecution(); + Saved next{history, {}}; + if (spec.rewind == Rewind::kSnapshot) next.snapshot = Save(); + saved = std::move(next); + return Position(); + } + + int Restore() { + Usable(); + if (!saved) return -1; + if (saved->history == history) return Position(); + auto lease = runtime->AcquireExecution(); + GoTo(saved->history); + return Position(); + } + + void Clear() { + Usable(); + // A seekable model never reads its cache past the position, so nothing needs clearing. + if (spec.rewind == Rewind::kSeek) { + history.clear(); + } else { + auto lease = runtime->AcquireExecution(); + ResetModel(); + } + saved.reset(); + } + + void UpdateMaxLength(std::uint32_t n) { + Usable(); + if (!n || n > spec.context_limit || n < static_cast(Position())) + throw std::out_of_range(spec.display_name + " maximum length is invalid"); + if (n > capacity) Grow(n); + if (saved && n < saved->history.size()) saved.reset(); + max_length = n; + } + +private: + void CheckIds(std::span ids) const { + for (const int id : ids) + if (id < 0 || static_cast(id) >= vocab) + throw std::out_of_range(spec.display_name + " token " + std::to_string(id) + + " is outside the vocabulary of " + std::to_string(vocab)); + } + + /// \brief run `fresh` after every token the model holds, and keep it + void Advance(std::span fresh, bool want_logits) { + const auto at = history.size(); + if (at + fresh.size() > max_length) + throw std::out_of_range(spec.display_name + + " request exceeds configured context capacity"); + const auto from = PassStart(at, fresh.size()); + std::vector rows(history.begin() + static_cast(from), + history.end()); + rows.insert(rows.end(), fresh.begin(), fresh.end()); + Dispatch(rows, from, want_logits); + history.insert(history.end(), fresh.begin(), fresh.end()); + } + + /// \brief where a pass adding `rows` tokens at position `at` has to start + std::size_t PassStart(std::size_t at, std::size_t rows) const { + if (at == 0 || rows == 1) return at; + switch (spec.continuation) { + case Continuation::kAnywhere: + return at; + case Continuation::kFromStart: + return 0; + case Continuation::kFromPage: + return capacity > kPageRows ? at - at % kPageRows : 0; + } + return 0; + } + + void Dispatch(const std::vector& rows, std::size_t from, bool want_logits) { + ryzenai_corelib_model_forward_options options{}; + options.position = + spec.rewind == Rewind::kSeek ? from : RYZENAI_CORELIB_POSITION_CONTINUE; + if (want_logits) { + options.logits = logits.data(); + options.logits_capacity = logits.size(); + } + try { + api->Check(api->functions().model_forward(model.get(), rows.data(), nullptr, + rows.size(), nullptr, &options), + "ryzenai_corelib_model_forward"); + } catch (...) { + Poison(); + throw; + } + } + + /// \brief make `target` the sequence the model holds, as cheaply as its state allows + void GoTo(std::vector target) { + const bool prefix = target.size() <= history.size() && + std::equal(target.begin(), target.end(), history.begin()); + if (spec.rewind == Rewind::kSeek && prefix) { + history = std::move(target); + return; + } + if (saved && saved->snapshot && saved->history == target) { + Load(saved->snapshot); + history = std::move(target); + return; + } + ResetModel(); + if (!target.empty()) Advance(target, false); + } + + /// \note The capacity is fixed for a model's life, so a larger one is a + /// new model over the same packed weights. + void Grow(std::uint32_t n) { + auto lease = runtime->AcquireExecution(); + UniqueModelSnapshot current; + if (spec.rewind == Rewind::kSnapshot && !history.empty()) current = Save(); + Recreate(n); + if (current) { + Load(current); + } else if (!history.empty()) { + auto replay = std::move(history); + history.clear(); + Advance(replay, false); + } + } + + void Recreate(std::uint32_t n) { + ryzenai_corelib_model_create_options options{}; + options.size = sizeof(options); + options.max_sequence_length = static_cast(n); + void* raw = nullptr; + api->Check(api->functions().model_create(spec.corelib_name.c_str(), weights.get(), + &options, &raw), + "ryzenai_corelib_model_create"); + UniqueModel created(api, raw); + ryzenai_corelib_model_info info{}; + api->Check(api->functions().model_info_get(created.get(), &info), + "ryzenai_corelib_model_info_get"); + if (info.vocab_size <= 0 || info.context_length < static_cast(n)) + throw std::runtime_error(spec.display_name + " model_info is vocab " + + std::to_string(info.vocab_size) + " context " + + std::to_string(info.context_length) + + "; this engine needs vocab > 0 and context >= " + + std::to_string(n)); + model = std::move(created); + capacity = static_cast(info.context_length); + vocab = static_cast(info.vocab_size); + logits.resize(vocab); + } + + UniqueModelSnapshot Save() { + void* raw = nullptr; + api->Check(api->functions().model_state_save(model.get(), &raw), + "ryzenai_corelib_model_state_save"); + return UniqueModelSnapshot(api, raw); + } + + /// \note corelib leaves the model reset when a restore fails partway, which + /// no history here describes. + void Load(const UniqueModelSnapshot& snapshot) { + try { + api->Check(api->functions().model_state_restore(model.get(), snapshot.get()), + "ryzenai_corelib_model_state_restore"); + } catch (...) { + Poison(); + throw; + } + } + + void ResetModel() { + try { + api->Check(api->functions().model_reset(model.get()), "ryzenai_corelib_model_reset"); + } catch (...) { + Poison(); + throw; + } + history.clear(); + } + + void Poison() { + poisoned = true; + history.clear(); + saved.reset(); + } + + buffer Logits() const { + buffer out(vocab); + F32ToBf16(logits, + std::span(reinterpret_cast(out.data()), vocab)); + return out; + } +}; + +CorelibModel::CorelibModel(ModelSpec spec, std::filesystem::path weights, + std::shared_ptr runtime, + std::uint32_t max_length) + : impl_(std::make_unique(std::move(spec), weights, std::move(runtime), max_length)) {} + +CorelibModel::~CorelibModel() = default; + +buffer CorelibModel::forward(int id) { return impl_->Forward(id); } + +buffer CorelibModel::prefill(std::vector& ids, void* payload) { + return impl_->Prefill(ids, payload); +} + +void CorelibModel::set_context_length(int length) { impl_->RewindTo(length); } + +// An ABI shim, not a capability: load_weights is pure virtual in the frozen +// causal_lm.hpp, and corelib packs this engine's weights from the file itself. +void CorelibModel::load_weights(Q4NX&) { + impl_->Usable(); + throw std::runtime_error(impl_->spec.display_name + " rai weights are loaded by corelib"); +} + +void CorelibModel::update_max_length(std::uint32_t max_length) { + impl_->UpdateMaxLength(max_length); +} + +void CorelibModel::clear_context() { impl_->Clear(); } + +buffer CorelibModel::get_k_cache(int, int) { + throw std::runtime_error(impl_->spec.display_name + " rai does not expose the KV cache"); +} + +buffer CorelibModel::get_v_cache(int, int) { + throw std::runtime_error(impl_->spec.display_name + " rai does not expose the KV cache"); +} + +int CorelibModel::get_current_context_length() { + impl_->Usable(); + return impl_->Position(); +} + +int CorelibModel::checkpoint() { return impl_->Checkpoint(); } + +int CorelibModel::restore() { return impl_->Restore(); } + +bool CorelibModel::poisoned() const noexcept { return impl_ && impl_->poisoned; } + +} // namespace flm::rai diff --git a/src/common/rai/corelib_model_backend.cpp b/src/common/rai/corelib_model_backend.cpp new file mode 100644 index 000000000..3452db6c9 --- /dev/null +++ b/src/common/rai/corelib_model_backend.cpp @@ -0,0 +1,205 @@ +/// \file corelib_model_backend.cpp +#include "rai/corelib_model_backend.hpp" + +#include "rai/corelib_runtime.hpp" +#include "utils/file_access.hpp" +#include "utils/utils.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace flm::rai { +namespace { + +nlohmann::json ReadJson(const std::filesystem::path& path) { + flm::file_access::ObserveOpen(path); + std::ifstream input(path, std::ios::binary); + if (!input) throw std::runtime_error("Cannot open " + path.string()); + try { + return nlohmann::json::parse(input); + } catch (const std::exception& error) { + throw std::runtime_error("Cannot parse " + path.string() + ": " + error.what()); + } +} + +bool NamesAGguf(const std::filesystem::path& path) { + auto extension = path.extension().string(); + std::transform(extension.begin(), extension.end(), extension.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + return extension == ".gguf"; +} + +/// \brief the catalog's one GGUF or, when it names none, the directory's one +std::filesystem::path RequireGguf(const nlohmann::ordered_json& model_info, + const std::filesystem::path& root, const std::string& name) { + std::vector named; + if (const auto files = model_info.find("files"); + files != model_info.end() && files->is_array()) { + for (const auto& file : *files) + if (file.is_string() && NamesAGguf(file.get())) + named.push_back(file.get()); + } + if (named.size() == 1) { + const auto path = root / named.front(); + if (!std::filesystem::exists(path)) + throw std::runtime_error(name + " rai: the catalog entry names " + named.front() + + " but " + path.string() + " does not exist"); + return path; + } + if (named.size() > 1) + throw std::runtime_error(name + " rai: the catalog entry names " + + std::to_string(named.size()) + + " GGUFs; it must name exactly one"); + std::vector found; + std::error_code failed; + for (const auto& entry : std::filesystem::directory_iterator(root, failed)) + if (entry.is_regular_file(failed) && NamesAGguf(entry.path())) found.push_back(entry.path()); + if (found.size() == 1) return found.front(); + throw std::runtime_error(name + " rai: the catalog entry names no GGUF and " + root.string() + + " holds " + std::to_string(found.size()) + + " of them; it must hold exactly one"); +} + +/// \brief a special token's text, whether stated bare or as {"content": ...} +std::optional TokenText(const nlohmann::json& config, const char* key) { + const auto it = config.find(key); + if (it == config.end()) return std::nullopt; + if (it->is_string()) return it->get(); + if (it->is_object() && it->contains("content") && (*it)["content"].is_string()) + return (*it)["content"].get(); + return std::nullopt; +} + +/// \brief the id tokenizer.json gives `text`, from its added tokens or its vocabulary +std::optional TokenId(const nlohmann::json& tokenizer, const std::string& text) { + if (const auto added = tokenizer.find("added_tokens"); + added != tokenizer.end() && added->is_array()) { + for (const auto& token : *added) + if (token.value("content", std::string()) == text && token.contains("id") && + token["id"].is_number_integer()) + return token["id"].get(); + } + if (const auto model = tokenizer.find("model"); + model != tokenizer.end() && model->is_object() && model->contains("vocab")) { + const auto& vocab = (*model)["vocab"]; + if (vocab.is_object() && vocab.contains(text) && vocab[text].is_number_integer()) + return vocab[text].get(); + } + return std::nullopt; +} + +struct PackageTokens { + std::vector eos; + std::optional bos; +}; + +PackageTokens ReadPackageTokens(const BackendSpec& spec, + const flm::backend::BackendContext& context) { + PackageTokens out; + out.eos = spec.eos_ids; + if (!spec.eos_ids.empty() && !spec.read_bos) return out; + + const std::filesystem::path root(context.model_path); + const auto& name = spec.model.display_name; + const auto config = context.tokenizer_config ? *context.tokenizer_config + : ReadJson(root / "tokenizer_config.json"); + const auto tokenizer = ReadJson(root / "tokenizer.json"); + const auto require = [&](const char* key) { + const auto text = TokenText(config, key); + if (!text) throw std::runtime_error(name + " tokenizer_config.json has no " + key); + const auto id = TokenId(tokenizer, *text); + if (!id) + throw std::runtime_error(name + " tokenizer.json has no id for " + key + " " + *text); + return *id; + }; + + if (spec.eos_ids.empty()) { + out.eos.push_back(require("eos_token")); + const auto model_config = root / "config.json"; + if (std::filesystem::exists(model_config)) { + const auto document = ReadJson(model_config); + if (!document.is_object()) + throw std::runtime_error(name + " " + model_config.string() + + " is not a JSON object"); + const auto listed = document.value("eos_token_id", nlohmann::json()); + std::vector ids; + if (listed.is_number_integer()) ids.push_back(listed.get()); + if (listed.is_array()) + for (const auto& id : listed) + if (id.is_number_integer()) ids.push_back(id.get()); + for (const int id : ids) + if (std::find(out.eos.begin(), out.eos.end(), id) == out.eos.end()) + out.eos.push_back(id); + } + } + if (spec.read_bos) out.bos = require("bos_token"); + return out; +} + +class CorelibBackend final : public flm::backend::ModelBackend { +public: + CorelibBackend(const BackendSpec& spec, const flm::backend::BackendContext& context) { + const auto& name = spec.model.display_name; + if (context.enable_preemption) + throw std::invalid_argument(name + " rai does not support preemption"); + const auto length = context.context_length; + const bool fits = spec.context_fits ? spec.context_fits(length) + : length >= 1 && length <= spec.model.context_limit; + if (!fits) + throw std::out_of_range(name + " rai context length must be in " + spec.context_rule); + context_length_ = length; + + auto tokens = ReadPackageTokens(spec, context); + eos_ids_ = std::move(tokens.eos); + bos_id_ = tokens.bos; + const auto gguf = RequireGguf(context.model_info, context.model_path, name); + + runtime_ = corelib::CorelibRuntime::GetOrCreate(utils::get_executable_directory()); + auto engine = std::make_unique(spec.model, gguf, runtime_, length); + engine->clear_context(); + engine_ = std::move(engine); + } + + ~CorelibBackend() override { engine_.reset(); } + + causal_lm& engine() override { return *engine_; } + std::string id() const override { return flm::backend::kRaiBackendId; } + std::string detail() const override { + return runtime_ ? runtime_->loaded_library_path().string() : std::string(); + } + std::uint32_t max_decode_length() const override { return context_length_ - 1; } + bool supports_preemption() const override { return false; } + /// \note corelib refuses a forward past its context, so the extra one the + /// FastFlowLM engines take after an EOS token would fail here. + bool forwards_past_eos() const override { return false; } + bool poisoned() const noexcept override { return engine_ && engine_->poisoned(); } + std::optional> forced_eos_ids() const override { return eos_ids_; } + std::optional forced_bos_id() const override { return bos_id_; } + +private: + std::uint32_t context_length_{}; + std::vector eos_ids_; + std::optional bos_id_; + // Declared before the engine so it is destroyed after it: corelib objects + // must not outlive the API they came from. + std::shared_ptr runtime_; + std::unique_ptr engine_; +}; + +} // namespace + +flm::backend::BackendFactory Factory(BackendSpec spec) { + return [spec = std::move(spec)](const flm::backend::BackendContext& context) + -> std::unique_ptr { + return std::make_unique(spec, context); + }; +} + +} // namespace flm::rai diff --git a/src/common/rai/gguf_file.cpp b/src/common/rai/gguf_file.cpp deleted file mode 100644 index c070cae4b..000000000 --- a/src/common/rai/gguf_file.cpp +++ /dev/null @@ -1,512 +0,0 @@ -#include "rai/gguf_file.hpp" - -#include "utils/file_access.hpp" - -#ifdef _WIN32 -#ifndef NOMINMAX -#define NOMINMAX -#endif -#include -#else -#include -#include -#include -#include -#include -#endif - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace flm::rai { -namespace { -constexpr std::uint32_t kMagic = 0x46554747; -constexpr std::uint32_t kVersion = 3; -constexpr std::uint64_t kDefaultAlignment = 32; - -std::uint64_t CheckedAdd(std::uint64_t a, std::uint64_t b, std::string_view field) { - if (a > std::numeric_limits::max() - b) - throw std::runtime_error(std::string(field) + ": overflow in addition"); - return a + b; -} - -std::uint64_t CheckedMultiply(std::uint64_t a, std::uint64_t b, std::string_view field) { - if (a != 0 && b > std::numeric_limits::max() / a) - throw std::runtime_error(std::string(field) + ": overflow in product"); - return a * b; -} - -std::uint64_t AlignUp(std::uint64_t value, std::uint64_t alignment) { - return CheckedAdd(value, alignment - 1, "alignment") & ~(alignment - 1); -} - -std::span RequireRange(std::span file, - std::uint64_t offset, std::uint64_t length, - std::string_view field) { - const auto end = CheckedAdd(offset, length, field); - if (end > file.size() || offset > std::numeric_limits::max() || - length > std::numeric_limits::max()) - FailField(field, "out-of-file range", "range within mapped file"); - return file.subspan(static_cast(offset), static_cast(length)); -} - -class Cursor { -public: - explicit Cursor(std::span file) : file_(file) {} - - template - T Read(std::string_view field) { - const auto bytes = RequireRange(file_, offset_, sizeof(T), field); - T value; - std::memcpy(&value, bytes.data(), sizeof(T)); - offset_ = CheckedAdd(offset_, sizeof(T), field); - return value; - } - - std::string ReadString(std::string_view field) { - const auto length = Read(field); - const auto bytes = RequireRange(file_, offset_, length, field); - std::string value(reinterpret_cast(bytes.data()), bytes.size()); - offset_ = CheckedAdd(offset_, length, field); - return value; - } - - void Skip(std::uint64_t length, std::string_view field) { - RequireRange(file_, offset_, length, field); - offset_ = CheckedAdd(offset_, length, field); - } - - std::uint64_t offset() const noexcept { return offset_; } - std::span remaining() const { - return file_.subspan(static_cast(offset_)); - } - -private: - std::span file_; - std::uint64_t offset_{}; -}; - -std::string MetadataTypeName(std::uint32_t type) { - static constexpr const char* names[] = {"UINT8", "INT8", "UINT16", "INT16", - "UINT32", "INT32", "FLOAT32", "BOOL", "STRING", "ARRAY", "UINT64", - "INT64", "FLOAT64"}; - return type < std::size(names) ? names[type] : "unknown(" + std::to_string(type) + ")"; -} - -std::uint64_t FixedMetadataSize(std::uint32_t type) { - switch (type) { - case 0: case 1: case 7: return 1; - case 2: case 3: return 2; - case 4: case 5: case 6: return 4; - case 10: case 11: case 12: return 8; - default: return 0; - } -} - -bool IsIntegerMetadata(std::uint32_t type) { - return type <= 5 || type == 10 || type == 11; -} - -/// \note `values` is filled only for a retained integer array; every other -/// array keeps just its count. -struct ArrayInfo { - std::uint32_t type; - std::uint64_t count; - std::vector values; - bool has_negative = false; -}; -using MetadataValue = std::variant; - -MetadataValue ReadMetadataValue(Cursor& cursor, std::uint32_t type, - std::string_view field, bool retain) { - switch (type) { - case 0: { auto v = cursor.Read(field); return retain ? MetadataValue(std::uint64_t(v)) : MetadataValue{}; } - case 1: { auto v = cursor.Read(field); return retain ? MetadataValue(std::int64_t(v)) : MetadataValue{}; } - case 2: { auto v = cursor.Read(field); return retain ? MetadataValue(std::uint64_t(v)) : MetadataValue{}; } - case 3: { auto v = cursor.Read(field); return retain ? MetadataValue(std::int64_t(v)) : MetadataValue{}; } - case 4: { auto v = cursor.Read(field); return retain ? MetadataValue(std::uint64_t(v)) : MetadataValue{}; } - case 5: { auto v = cursor.Read(field); return retain ? MetadataValue(std::int64_t(v)) : MetadataValue{}; } - case 6: { auto v = cursor.Read(field); return retain ? MetadataValue(double(v)) : MetadataValue{}; } - case 7: { auto v = cursor.Read(field); if (v > 1) FailField(field, std::to_string(v), "GGUF boolean 0 or 1"); return retain ? MetadataValue(bool(v)) : MetadataValue{}; } - case 8: { auto v = cursor.ReadString(field); return retain ? MetadataValue(std::move(v)) : MetadataValue{}; } - case 9: { - const std::string array_field = std::string(field) + " array"; - const auto element_type = cursor.Read(array_field); - const auto count = cursor.Read(array_field); - if (element_type == 9 || element_type > 12) - FailField(array_field, MetadataTypeName(element_type), "a skippable GGUF array element type"); - ArrayInfo info{element_type, count}; - const auto fixed = FixedMetadataSize(element_type); - if (retain && IsIntegerMetadata(element_type)) { - const auto bytes = CheckedMultiply(count, fixed, array_field); - RequireRange(cursor.remaining(), 0, bytes, array_field); - info.values.reserve(static_cast(count)); - for (std::uint64_t i = 0; i < count; ++i) { - const auto element = ReadMetadataValue(cursor, element_type, array_field, true); - if (const auto* u = std::get_if(&element)) { - info.values.push_back(*u); - } else { - const auto s = std::get(element); - if (s < 0) info.has_negative = true; - info.values.push_back(static_cast(s)); - } - } - } else if (fixed != 0) { - cursor.Skip(CheckedMultiply(count, fixed, array_field), array_field); - } else { - for (std::uint64_t i = 0; i < count; ++i) - (void)ReadMetadataValue(cursor, element_type, array_field, false); - } - return retain ? MetadataValue(std::move(info)) : MetadataValue{}; - } - case 10: { auto v = cursor.Read(field); return retain ? MetadataValue(v) : MetadataValue{}; } - case 11: { auto v = cursor.Read(field); return retain ? MetadataValue(v) : MetadataValue{}; } - case 12: { auto v = cursor.Read(field); return retain ? MetadataValue(v) : MetadataValue{}; } - default: - FailField(field, MetadataTypeName(type), "a supported metadata type"); - } -} - -std::string ShapeText(std::span shape) { - std::ostringstream out; - out << '['; - for (std::size_t i = 0; i < shape.size(); ++i) { - if (i) out << ','; - out << shape[i]; - } - return out.str() + ']'; -} - -std::string GgmlTypeName(std::uint32_t type) { - if (type == kGgmlTypeF32) return "F32"; - if (type == kGgmlTypeQ8_0) return "Q8_0"; - if (type == kGgmlTypeBF16) return "BF16"; - return "GGML type " + std::to_string(type); -} - -std::uint64_t ElementCount(std::span shape, std::string_view field) { - std::uint64_t result = 1; - for (const auto dimension : shape) { - if (dimension <= 0) FailField(field, std::to_string(dimension), "positive dimensions"); - result = CheckedMultiply(result, static_cast(dimension), field); - } - return result; -} - -/// \return the byte length, or 0 for a type this reader does not size -/// \note A file may carry tensors of types no rai engine reads (a K-quant -/// output head, say); those are indexed without a range check and -/// rejected only if something asks for them. -std::uint64_t TensorByteLength(std::uint32_t type, std::span shape, - std::string_view field) { - const auto elements = ElementCount(shape, field); - if (type == kGgmlTypeF32) return CheckedMultiply(elements, 4, field); - if (type == kGgmlTypeBF16) return CheckedMultiply(elements, 2, field); - if (type == kGgmlTypeQ8_0) { - if (elements % 32 != 0) - FailField(field, std::to_string(elements) + " elements", "Q8_0 element count divisible by 32"); - return CheckedMultiply(elements / 32, 34, field); - } - return 0; -} -} // namespace - -void FailField(std::string_view field, std::string actual, std::string expected) { - throw std::runtime_error(std::string(field) + ": actual " + actual + - ", expected " + expected); -} - -struct GgufFile::Impl { - struct TensorRecord { - std::string name; - std::span bytes; - std::vector shape; - std::uint32_t type; - std::uint64_t absolute_offset; - }; - - std::filesystem::path path; -#ifdef _WIN32 - HANDLE file = INVALID_HANDLE_VALUE; - HANDLE mapping = nullptr; -#else - int file = -1; -#endif - const std::byte* data = nullptr; - std::uint64_t size = 0; - std::map> tensors; - std::map> metadata; - std::map> metadata_types; - - ~Impl() { -#ifdef _WIN32 - if (data) UnmapViewOfFile(data); - if (mapping) CloseHandle(mapping); - if (file != INVALID_HANDLE_VALUE) CloseHandle(file); -#else - if (data) - ::munmap(const_cast(data), static_cast(size)); - if (file >= 0) ::close(file); -#endif - } - - std::span bytes() const { return {data, static_cast(size)}; } - - const TensorRecord& Record(std::string_view name) const { - const auto it = tensors.find(name); - if (it == tensors.end()) FailField(name, "missing", "present tensor"); - return it->second; - } - - const MetadataValue& Value(std::string_view key, const char* expected) const { - const auto it = metadata.find(key); - if (it == metadata.end()) FailField(key, "missing", expected); - return it->second; - } - - [[noreturn]] void WrongType(std::string_view key, const char* expected) const { - FailField(key, MetadataTypeName(metadata_types.find(key)->second), expected); - } - - void CheckShape(const TensorRecord& tensor, std::uint32_t type, - std::span expected) const { - if (tensor.type != type) FailField(tensor.name, GgmlTypeName(tensor.type), GgmlTypeName(type)); - if (!std::equal(tensor.shape.begin(), tensor.shape.end(), expected.begin(), expected.end())) - FailField(tensor.name, ShapeText(tensor.shape), ShapeText(expected)); - const auto length = TensorByteLength(type, expected, tensor.name); - if (tensor.bytes.size() != length) - FailField(tensor.name, std::to_string(tensor.bytes.size()) + " bytes", - std::to_string(length) + " bytes"); - } -}; - -GgufFile::GgufFile(std::unique_ptr impl) : impl_(std::move(impl)) {} -GgufFile::~GgufFile() = default; - -std::shared_ptr GgufFile::Open(const std::filesystem::path& path, - const RetainKey& retain) { - auto impl = std::make_unique(); - impl->path = path; - flm::file_access::ObserveOpen(path); -#ifdef _WIN32 - impl->file = CreateFileW(path.c_str(), GENERIC_READ, FILE_SHARE_READ, nullptr, - OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); - if (impl->file == INVALID_HANDLE_VALUE) - throw std::runtime_error("GGUF file " + path.string() + ": actual open failure " + - std::to_string(GetLastError()) + ", expected readable file"); - LARGE_INTEGER size; - if (!GetFileSizeEx(impl->file, &size) || size.QuadPart <= 0 || - static_cast(size.QuadPart) > std::numeric_limits::max()) - FailField("GGUF file size", std::to_string(size.QuadPart), "positive mappable size"); - impl->size = static_cast(size.QuadPart); - impl->mapping = CreateFileMappingW(impl->file, nullptr, PAGE_READONLY, 0, 0, nullptr); - if (!impl->mapping) - throw std::runtime_error("GGUF mapping: actual CreateFileMappingW failure " + - std::to_string(GetLastError()) + ", expected PAGE_READONLY mapping"); - impl->data = static_cast(MapViewOfFile(impl->mapping, FILE_MAP_READ, 0, 0, 0)); - if (!impl->data) - throw std::runtime_error("GGUF mapping: actual MapViewOfFile failure " + - std::to_string(GetLastError()) + ", expected FILE_MAP_READ view"); -#else - impl->file = ::open(path.c_str(), O_RDONLY | O_CLOEXEC); - if (impl->file < 0) - throw std::runtime_error("GGUF file " + path.string() + ": actual open failure " + - std::to_string(errno) + ", expected readable file"); - struct ::stat status {}; - if (::fstat(impl->file, &status) != 0 || status.st_size <= 0 || - static_cast(status.st_size) > - std::numeric_limits::max()) - FailField("GGUF file size", - std::to_string(static_cast(status.st_size)), - "positive mappable size"); - impl->size = static_cast(status.st_size); - void* view = ::mmap(nullptr, static_cast(impl->size), PROT_READ, - MAP_PRIVATE, impl->file, 0); - if (view == MAP_FAILED) - throw std::runtime_error("GGUF mapping: actual mmap failure " + - std::to_string(errno) + ", expected PROT_READ view"); - impl->data = static_cast(view); -#endif - - const auto file = impl->bytes(); - Cursor cursor(file); - if (cursor.Read("GGUF header") != kMagic) - FailField("GGUF magic", "mismatch", "0x46554747"); - const auto version = cursor.Read("GGUF header"); - if (version != kVersion) FailField("GGUF version", std::to_string(version), "3"); - const auto tensor_count = cursor.Read("tensor count"); - const auto metadata_count = cursor.Read("metadata count"); - if (tensor_count > file.size() / 24) - FailField("tensor count", std::to_string(tensor_count), "count fitting directory"); - if (metadata_count > file.size() / 12) - FailField("metadata count", std::to_string(metadata_count), "count fitting metadata"); - - for (std::uint64_t i = 0; i < metadata_count; ++i) { - const auto key = cursor.ReadString("metadata key string"); - const auto type = cursor.Read(key); - const bool keep = key == "general.alignment" || !retain || retain(key); - auto value = ReadMetadataValue(cursor, type, key, keep); - if (keep) { - if (!impl->metadata.emplace(key, std::move(value)).second) - FailField(key, "duplicate metadata key", "unique metadata key"); - impl->metadata_types.emplace(key, type); - } - } - - std::uint64_t alignment = kDefaultAlignment; - if (impl->metadata.contains("general.alignment")) { - const auto* value = std::get_if(&impl->metadata.at("general.alignment")); - if (!value) impl->WrongType("general.alignment", "unsigned integer metadata"); - alignment = *value; - } - if (alignment == 0 || (alignment & (alignment - 1)) != 0) - FailField("general.alignment", std::to_string(alignment), "a non-zero power of two"); - - struct DirectoryTensor { - std::string name; - std::vector shape; - std::uint32_t type; - std::uint64_t relative_offset; - std::uint64_t length; - }; - std::vector directory; - directory.reserve(static_cast(tensor_count)); - for (std::uint64_t i = 0; i < tensor_count; ++i) { - auto name = cursor.ReadString("tensor directory name"); - const auto dimension_count = cursor.Read("tensor directory dimensions"); - if (dimension_count == 0 || dimension_count > 4) - FailField(name, std::to_string(dimension_count), "1..4 tensor dimensions"); - std::vector shape; - shape.reserve(dimension_count); - for (std::uint32_t d = 0; d < dimension_count; ++d) { - const auto dimension = cursor.Read("tensor directory dimension"); - if (dimension > static_cast(std::numeric_limits::max())) - FailField(name, std::to_string(dimension), "dimension fitting int64"); - shape.push_back(static_cast(dimension)); - } - std::reverse(shape.begin(), shape.end()); - const auto type = cursor.Read("tensor directory type"); - const auto offset = cursor.Read("tensor directory offset"); - const auto length = TensorByteLength(type, shape, name); - directory.push_back({std::move(name), std::move(shape), type, offset, length}); - } - - const auto data_start = AlignUp(cursor.offset(), alignment); - struct Range { std::uint64_t begin, end; std::string name; }; - std::vector ranges; - ranges.reserve(directory.size()); - for (auto& tensor : directory) { - if (tensor.relative_offset % alignment != 0) - FailField(tensor.name, std::to_string(tensor.relative_offset), - "offset aligned to " + std::to_string(alignment)); - const auto absolute = CheckedAdd(data_start, tensor.relative_offset, tensor.name); - std::span bytes; - if (tensor.length != 0) { - bytes = RequireRange(file, absolute, tensor.length, tensor.name + " range"); - ranges.push_back({absolute, CheckedAdd(absolute, tensor.length, tensor.name), tensor.name}); - } - auto [it, inserted] = impl->tensors.emplace(tensor.name, - Impl::TensorRecord{tensor.name, bytes, std::move(tensor.shape), tensor.type, absolute}); - if (!inserted) FailField(tensor.name, "duplicate tensor name", "unique tensor name"); - } - std::sort(ranges.begin(), ranges.end(), [](const Range& a, const Range& b) { return a.begin < b.begin; }); - for (std::size_t i = 1; i < ranges.size(); ++i) { - if (ranges[i].begin < ranges[i - 1].end) - FailField(ranges[i].name, "overlap with " + ranges[i - 1].name, "non-overlapping tensor range"); - } - return std::shared_ptr(new GgufFile(std::move(impl))); -} - -const std::filesystem::path& GgufFile::Path() const { return impl_->path; } - -bool GgufFile::HasTensor(std::string_view name) const { return impl_->tensors.contains(name); } - -bool GgufFile::HasMetadata(std::string_view key) const { return impl_->metadata.contains(key); } - -GgufTensor GgufFile::RequireQ8(std::string_view name, std::span shape) const { - const auto& tensor = impl_->Record(name); - impl_->CheckShape(tensor, kGgmlTypeQ8_0, shape); - return {tensor.name, tensor.bytes, tensor.shape, tensor.type}; -} - -GgufFloatTensor GgufFile::RequireF32(std::string_view name, std::span shape) const { - const auto& tensor = impl_->Record(name); - impl_->CheckShape(tensor, kGgmlTypeF32, shape); - const auto address = reinterpret_cast(tensor.bytes.data()); - if (tensor.absolute_offset % alignof(float) != 0 || address % alignof(float) != 0) - FailField(name, "address/offset not aligned", "alignment 4"); - return {tensor.name, - {reinterpret_cast(tensor.bytes.data()), tensor.bytes.size() / sizeof(float)}, - tensor.shape}; -} - -GgufBf16Tensor GgufFile::RequireBf16(std::string_view name, std::span shape) const { - const auto& tensor = impl_->Record(name); - impl_->CheckShape(tensor, kGgmlTypeBF16, shape); - const auto address = reinterpret_cast(tensor.bytes.data()); - if (tensor.absolute_offset % alignof(std::uint16_t) != 0 || address % alignof(std::uint16_t) != 0) - FailField(name, "address/offset not aligned", "alignment 2"); - return {tensor.name, - {reinterpret_cast(tensor.bytes.data()), - tensor.bytes.size() / sizeof(std::uint16_t)}, - tensor.shape, tensor.type}; -} - -GgufTensor GgufFile::Tensor(std::string_view name) const { - const auto& tensor = impl_->Record(name); - return {tensor.name, tensor.bytes, tensor.shape, tensor.type}; -} - -std::size_t GgufFile::TensorCount() const noexcept { return impl_->tensors.size(); } - -std::uint64_t GgufFile::Unsigned(std::string_view key) const { - const auto& value = impl_->Value(key, "unsigned integer metadata"); - if (const auto* v = std::get_if(&value)) return *v; - if (const auto* v = std::get_if(&value); v && *v >= 0) - return static_cast(*v); - impl_->WrongType(key, "unsigned integer metadata"); -} - -double GgufFile::Number(std::string_view key) const { - const auto& value = impl_->Value(key, "floating-point metadata"); - if (const auto* v = std::get_if(&value)) return *v; - impl_->WrongType(key, "floating-point metadata"); -} - -bool GgufFile::Boolean(std::string_view key) const { - const auto& value = impl_->Value(key, "boolean metadata"); - if (const auto* v = std::get_if(&value)) return *v; - impl_->WrongType(key, "boolean metadata"); -} - -std::string GgufFile::String(std::string_view key) const { - const auto& value = impl_->Value(key, "string metadata"); - if (const auto* v = std::get_if(&value)) return *v; - impl_->WrongType(key, "string metadata"); -} - -std::uint64_t GgufFile::ArrayCount(std::string_view key) const { - const auto& value = impl_->Value(key, "array metadata"); - if (const auto* v = std::get_if(&value)) return v->count; - impl_->WrongType(key, "array metadata"); -} - -std::vector GgufFile::UnsignedArray(std::string_view key) const { - const auto& value = impl_->Value(key, "unsigned integer array metadata"); - const auto* array = std::get_if(&value); - if (!array) impl_->WrongType(key, "unsigned integer array metadata"); - if (!IsIntegerMetadata(array->type)) - FailField(key, "array of " + MetadataTypeName(array->type), "array of integers"); - if (array->has_negative) - FailField(key, "negative array element", "non-negative integers"); - return array->values; -} - -} // namespace flm::rai diff --git a/src/common/rai/host_ops.cpp b/src/common/rai/host_ops.cpp index a08aacadb..19324c471 100644 --- a/src/common/rai/host_ops.cpp +++ b/src/common/rai/host_ops.cpp @@ -1,100 +1,9 @@ #include "rai/host_ops.hpp" #include -#include -#include -#include #include namespace flm::rai { -namespace { - -float HalfToFloat(std::uint16_t half) { - const std::uint32_t sign = static_cast(half & 0x8000) << 16; - const std::uint32_t exponent = (half >> 10) & 0x1f; - std::uint32_t fraction = half & 0x03ff; - std::uint32_t bits; - if (exponent == 0) { - if (fraction == 0) { - bits = sign; - } else { - int shift = 0; - while ((fraction & 0x0400) == 0) { - fraction <<= 1; - ++shift; - } - fraction &= 0x03ff; - bits = sign | (static_cast(127 - 14 - shift) << 23) | (fraction << 13); - } - } else if (exponent == 0x1f) { - bits = sign | 0x7f800000 | (fraction << 13); - } else { - bits = sign | ((exponent + (127 - 15)) << 23) | (fraction << 13); - } - return std::bit_cast(bits); -} - -} // namespace - -std::vector DecodeQ8Rows(const GgufTensor& table, std::span ids) { - if (table.ggml_type != kGgmlTypeQ8_0 || table.shape.size() != 2 || - table.shape[0] <= 0 || table.shape[1] <= 0 || table.shape[1] % 32 != 0) - throw std::runtime_error("embedding must be a two-dimensional Q8_0 tensor with block-aligned rows"); - const auto rows = static_cast(table.shape[0]); - const auto width = static_cast(table.shape[1]); - const auto blocks_per_row = width / 32; - const auto row_bytes = blocks_per_row * 34; - if (rows > std::numeric_limits::max() / row_bytes || - table.bytes.size() != rows * row_bytes) - throw std::runtime_error("embedding Q8_0 byte length does not match its logical shape"); - - std::vector result; - result.reserve(ids.size() * width); - for (const int id : ids) { - if (id < 0 || static_cast(id) >= rows) - throw std::out_of_range("embedding token id is outside the vocabulary"); - const std::byte* row = table.bytes.data() + static_cast(id) * row_bytes; - for (std::size_t block = 0; block < blocks_per_row; ++block) { - const std::byte* encoded = row + block * 34; - std::uint16_t scale_bits; - std::memcpy(&scale_bits, encoded, sizeof(scale_bits)); - const float scale = HalfToFloat(scale_bits); - for (std::size_t element = 0; element < 32; ++element) { - const auto code = static_cast(std::to_integer(encoded[2 + element])); - result.push_back(scale * static_cast(code)); - } - } - } - return result; -} - -void RmsNorm(std::span input, std::span scale, - std::int64_t rows, std::int64_t width, float epsilon, - std::span output) { - if (rows <= 0 || width <= 0) throw std::invalid_argument("RMSNorm rows and width must be positive"); - const auto row_count = static_cast(rows); - const auto row_width = static_cast(width); - if (row_count > std::numeric_limits::max() / row_width) - throw std::invalid_argument("RMSNorm shape overflow"); - const auto elements = row_count * row_width; - if (input.size() != elements || output.size() != elements || scale.size() != row_width) - throw std::invalid_argument("RMSNorm shape mismatch"); - if (!std::isfinite(epsilon) || epsilon < 0.0f) - throw std::invalid_argument("RMSNorm epsilon must be finite and nonnegative"); - - for (std::size_t row = 0; row < row_count; ++row) { - const auto base = row * row_width; - double sum_of_squares = 0.0; - for (std::size_t column = 0; column < row_width; ++column) { - const double value = input[base + column]; - sum_of_squares += value * value; - } - const float mean_square = static_cast(sum_of_squares / static_cast(width)); - const float denominator = std::sqrt(mean_square + epsilon); - for (std::size_t column = 0; column < row_width; ++column) - output[base + column] = (input[base + column] / denominator) * scale[column]; - } -} std::vector F32ToBf16(std::span values) { std::vector result(values.size()); diff --git a/src/common/rai/rai_sources.cmake b/src/common/rai/rai_sources.cmake index b2ea13a7b..80b26ab27 100644 --- a/src/common/rai/rai_sources.cmake +++ b/src/common/rai/rai_sources.cmake @@ -2,8 +2,8 @@ include("${CMAKE_CURRENT_LIST_DIR}/../models/models_sources.cmake") set(FLM_RAI_SOURCES "${CMAKE_CURRENT_LIST_DIR}/corelib_api.cpp" + "${CMAKE_CURRENT_LIST_DIR}/corelib_model.cpp" + "${CMAKE_CURRENT_LIST_DIR}/corelib_model_backend.cpp" "${CMAKE_CURRENT_LIST_DIR}/corelib_runtime.cpp" - "${CMAKE_CURRENT_LIST_DIR}/gguf_file.cpp" "${CMAKE_CURRENT_LIST_DIR}/host_ops.cpp" - "${CMAKE_CURRENT_LIST_DIR}/weight_cache.cpp" ${FLM_MODELS_RAI_SOURCES}) diff --git a/src/common/rai/weight_cache.cpp b/src/common/rai/weight_cache.cpp deleted file mode 100644 index 0427dca91..000000000 --- a/src/common/rai/weight_cache.cpp +++ /dev/null @@ -1,149 +0,0 @@ -#include "rai/weight_cache.hpp" - -#include - -#include -#include -#include - -namespace flm::rai { -namespace { - -/// bumped when the on-disk layout changes in a way older indexes cannot express -constexpr int kIndexFormat = 1; - -bool IsDisabled(std::string_view value) { - return value == "0" || value == "off" || value == "OFF" || value == "false"; -} - -} // namespace - -WeightCacheKey MakeWeightCacheKey(const std::filesystem::path& gguf_path, - std::uint32_t corelib_major, - std::uint32_t corelib_minor, - std::uint32_t corelib_patch, - std::string layout, - std::uint64_t weight_count) { - WeightCacheKey key; - key.corelib_major = corelib_major; - key.corelib_minor = corelib_minor; - key.corelib_patch = corelib_patch; - key.layout = std::move(layout); - key.weight_count = weight_count; - std::error_code error; - const auto size = std::filesystem::file_size(gguf_path, error); - if (!error) key.gguf_size = size; - const auto written = std::filesystem::last_write_time(gguf_path, error); - if (!error) key.gguf_write_time = written.time_since_epoch().count(); - return key; -} - -std::optional WeightCache::ForGguf(const std::filesystem::path& gguf_path) { - const auto stem = gguf_path.stem().string() + ".rai-weights"; - const char* configured = std::getenv("FLM_RAI_WEIGHT_CACHE"); - if (configured && *configured) { - if (IsDisabled(configured)) return std::nullopt; - return WeightCache(configured, stem); - } - return WeightCache(gguf_path.parent_path(), stem); -} - -WeightCache::WeightCache(std::filesystem::path directory, std::string stem) - : directory_(std::move(directory)), stem_(std::move(stem)) {} - -std::filesystem::path WeightCache::DataPath() const { return directory_ / (stem_ + ".bin"); } - -std::filesystem::path WeightCache::IndexPath() const { return directory_ / (stem_ + ".json"); } - -std::uint64_t WeightCache::Remove() const { - std::uint64_t reclaimed = 0; - const auto data = DataPath().string(); - const auto index = IndexPath().string(); - for (const auto& name : {data, data + ".tmp", index, index + ".tmp"}) { - std::error_code error; - const auto size = std::filesystem::file_size(name, error); - if (error) continue; - if (std::filesystem::remove(name, error) && !error) reclaimed += size; - } - return reclaimed; -} - -std::optional> WeightCache::ReadIndex( - const WeightCacheKey& expected) const { - try { - std::ifstream input(IndexPath(), std::ios::binary); - if (!input) return std::nullopt; - const auto document = nlohmann::json::parse(input, nullptr, false); - if (document.is_discarded() || document.value("format", 0) != kIndexFormat) - return std::nullopt; - - WeightCacheKey key; - key.gguf_size = document.value("gguf_size", std::uint64_t{0}); - key.gguf_write_time = document.value("gguf_write_time", std::int64_t{0}); - key.corelib_major = document.value("corelib_major", std::uint32_t{0}); - key.corelib_minor = document.value("corelib_minor", std::uint32_t{0}); - key.corelib_patch = document.value("corelib_patch", std::uint32_t{0}); - key.layout = document.value("layout", std::string()); - key.weight_count = document.value("weight_count", std::uint64_t{0}); - if (!(key == expected)) return std::nullopt; - - const auto spans = document.find("spans"); - if (spans == document.end() || !spans->is_array() || spans->size() != expected.weight_count) - return std::nullopt; - std::error_code error; - const auto data_size = std::filesystem::file_size(DataPath(), error); - if (error) return std::nullopt; - std::vector result; - result.reserve(spans->size()); - for (const auto& span : *spans) { - if (!span.is_object()) return std::nullopt; - CachedWeightSpan entry{span.value("offset", std::uint64_t{0}), - span.value("size", std::uint64_t{0})}; - // A span past the end means the index describes a file that was - // truncated under it. - if (entry.size == 0 || entry.offset + entry.size > data_size) return std::nullopt; - result.push_back(entry); - } - return result; - } catch (...) { - return std::nullopt; - } -} - -bool WeightCache::WriteIndex(const WeightCacheKey& key, - const std::vector& spans) const { - try { - nlohmann::json document; - document["format"] = kIndexFormat; - document["gguf_size"] = key.gguf_size; - document["gguf_write_time"] = key.gguf_write_time; - document["corelib_major"] = key.corelib_major; - document["corelib_minor"] = key.corelib_minor; - document["corelib_patch"] = key.corelib_patch; - document["layout"] = key.layout; - document["weight_count"] = key.weight_count; - auto array = nlohmann::json::array(); - for (const auto& span : spans) array.push_back({{"offset", span.offset}, {"size", span.size}}); - document["spans"] = std::move(array); - - const auto final_path = IndexPath(); - const auto temporary = final_path.string() + ".tmp"; - { - std::ofstream output(temporary, std::ios::binary | std::ios::trunc); - if (!output) return false; - output << document.dump(); - if (!output) return false; - } - std::error_code error; - std::filesystem::rename(temporary, final_path, error); - if (error) { - std::filesystem::remove(temporary, error); - return false; - } - return true; - } catch (...) { - return false; - } -} - -} // namespace flm::rai diff --git a/src/include/models/gemma4/rai/aie_next/gemma4_rai.hpp b/src/include/models/gemma4/rai/aie_next/gemma4_rai.hpp deleted file mode 100644 index c6c8d878c..000000000 --- a/src/include/models/gemma4/rai/aie_next/gemma4_rai.hpp +++ /dev/null @@ -1,126 +0,0 @@ -#pragma once - -#include "causal_lm.hpp" -#include "lm_config.hpp" -#include "models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp" -#include "rai/corelib_runtime.hpp" - -#include -#include -#include -#include -#include - -namespace flm::gemma4 { - -/// \brief one finished packed weight, and the named slot it was placed in -struct WeightPlacement { - /// \brief the GGUF tensor this slot holds, e.g. "blk.7.attn_q_norm.weight"; - /// "blk.7.ssmlp" and "blk.7.ple" for the two fused blocks - std::string slot; - const void* handle; -}; - -/// \brief the Gemma 4 engine on ryzenai-corelib: every packed weight the model -/// needs, held for the life of the object -/// -/// Construction opens one stream on the (8, 17) PDI pair, builds the shape -/// plan, creates every packed weight (matmul, ssmlp, rmsnorm and ple), -/// allocates the activation buffers, the non-uniform KV caches and both -/// rotary table pairs. `forward` and `prefill` run the layer loop and return -/// softcapped logits for the last position. -/// -/// A layer that owns its KV cache (`index < Gemma4Config::kv_layers`) has ten -/// weight objects; a layer that shares an earlier layer's cache has six -- it -/// never projects or norms K and V, although the file still carries those -/// tensors. -class gemma4_rai final : public causal_lm { -public: - /// \param package the GGUF and the config derived from it - /// \param runtime the corelib runtime this engine's objects belong to - /// \param config the frontend's model configuration; not read for shape, - /// every dimension comes from the GGUF - /// \param max_length the context this engine is built for: 1..4096, or a - /// multiple of 4096 up to kMaxPagedContext, and no longer than the - /// GGUF's own `gemma4.context_length` - gemma4_rai(std::shared_ptr package, - std::shared_ptr runtime, - const LM_Config& config, - std::uint32_t max_length = 4096); - - /// \brief construct with an overridden shape, for tests only - /// \note The shipped rows use one group size everywhere, so only an - /// overridden shape can tell `group`, `head_group` and `ple_group` - /// apart. `shape` is not validated against the file. - static std::unique_ptr CreateWithShapeForTest( - std::shared_ptr package, - std::shared_ptr runtime, - const LM_Config& config, Gemma4Config shape, - std::uint32_t max_length = 4096); - - ~gemma4_rai() override; - - /// \brief one decode step: one token in, softcapped logits out - /// \throws std::out_of_range if the token is outside the vocabulary, or - /// the request runs past the configured context - buffer forward(int id) override; - /// \brief one multi-token pass: `ids` in, softcapped logits for the last - /// position out - /// \note Same path as `forward`. A multi-row pass must start at position 0 - /// or, on a paged cache, on a page boundary, so a multi-row request - /// at any other position re-runs the tokens behind it from the - /// nearest such start together with `ids`: from position 0 unpaged, - /// from the current page's first row paged. - /// \throws std::invalid_argument for an empty request, or a continuation - /// whose earlier tokens this engine does not hold - /// \throws std::out_of_range as `forward` - buffer prefill(std::vector& ids, void* payload = nullptr) override; - void set_context_length(int length) override; - /// \brief unsupported; weights come from the GGUF package. Present because - /// causal_lm.hpp is a frozen ABI. - /// \throws std::runtime_error always - void load_weights(Q4NX&) override; - void update_max_length(std::uint32_t max_length) override; - void clear_context() override; - /// \brief one position's K row, gathered across the KV heads - /// \note Resolved through `Gemma4Config::layer_cache_owner`: a sharing - /// layer returns the cache it attends over. The width is - /// `kv_heads * head` of the owner, so it varies across layers. - /// \throws std::out_of_range if the layer or the index is out of range - buffer get_k_cache(int layer, int index) override; - /// \brief one position's V row; see get_k_cache - buffer get_v_cache(int layer, int index) override; - int get_current_context_length() override; - int checkpoint() override; - int restore() override; - - /// \brief whether a forward pass failed after a dispatch was accepted - /// \note Set when a pass throws after corelib accepted one of its - /// dispatches, or when its result is non-finite. Requests refused - /// before that (empty, bad token, chunked prefill, past capacity) do - /// not poison. Nothing clears it; every other entry point throws - /// while it is set, and the only exit is unload and reload. The - /// three const accessors below stay usable. - bool poisoned() const noexcept; - - /// \brief whether this load bound packed weights from the on-disk cache - bool loaded_from_cache() const noexcept; - - /// \brief how many packed weight objects this engine holds; also the - /// weight-cache key's slot count - std::size_t weight_slot_count() const noexcept; - - /// \brief every packed weight paired with the slot it was placed in, for - /// tests only - /// \note Reads the named members, not the slot table, so that it is an - /// independent statement of which tensor belongs in which member. - /// Size is `weight_slot_count()`. - std::vector WeightPlacementsForTest() const; - -private: - struct Impl; - explicit gemma4_rai(std::unique_ptr impl); - std::unique_ptr impl_; -}; - -} // namespace flm::gemma4 diff --git a/src/include/models/gemma4/rai/aie_next/gemma4_rai_backend.hpp b/src/include/models/gemma4/rai/aie_next/gemma4_rai_backend.hpp index 31cd48167..768c549e8 100644 --- a/src/include/models/gemma4/rai/aie_next/gemma4_rai_backend.hpp +++ b/src/include/models/gemma4/rai/aie_next/gemma4_rai_backend.hpp @@ -1,49 +1,41 @@ /// \file gemma4_rai_backend.hpp /// \brief The ryzenai-corelib backend for Gemma 4 (E2B and E4B) -/// \note Compiled only when FLM_ENABLE_RAI is on. A ModelBackend owns one -/// engine AND every rule for driving it, so everything that would -/// otherwise have to become a rai special case inside the gemma4e -/// frontend -- the 4095-token decode window, the no-preemption rule, the -/// package-verified EOS ids, the poisoning state -- is stated here -/// instead. The frontend stays backend-agnostic: there is no -/// `#if FLM_ENABLE_RAI` anywhere under src/common/AutoModel/ except -/// builtin_backends.cpp, which is the one file that knows both a family -/// name and an engine type. +/// \note Compiled only when FLM_ENABLE_RAI is on. #pragma once #include "AutoModel/model_backend.hpp" -#include "models/gemma4/rai/aie_next/gemma4_rai_constants.hpp" #include namespace flm::gemma4 { -/// \brief the largest context this backend can hold -/// \note THE ENGINE'S OWN CONSTANT. A context in 1..4096 stays on flat_mha. -/// Above that the engine pages at kPageSize, up to kMaxPagedContext. -inline constexpr std::uint32_t kRaiContextLimit = - static_cast(kMaxPagedContext); +/// \brief rows in one page of the paged KV cache, and the longest unpaged context +inline constexpr std::uint32_t kPageSize = 4096; + +/// \brief the longest context validated on hardware +/// \note corelib pages Gemma 4 up to the model's own 131072 and refuses past a +/// GGUF's smaller gemma4.context_length; 16384 is what was validated. +inline constexpr std::uint32_t kRaiContextLimit = 16384; /// \brief one below the longest context: the last position has to be readable -inline constexpr std::uint32_t kRaiDecodeLimit = - static_cast(kMaxPagedContext - 1); - -/// \brief what the frontend must know before building this backend -/// \return no xclbin, no preemption, 16384-token context ceiling -/// \note Pure data, and inline on purpose: `AutoModel::_shared_load_backend` -/// reads it while assembling the BackendContext -- that is, BEFORE any -/// engine exists -- so it cannot be a virtual on ModelBackend. Being in -/// the header also lets a test read it without linking the engine. +inline constexpr std::uint32_t kRaiDecodeLimit = kRaiContextLimit - 1; + +/// \return true when `length` is 1..4096, or a multiple of kPageSize up to kRaiContextLimit +inline bool ContextFits(std::uint32_t length) { + if (length < 1 || length > kRaiContextLimit) return false; + return length <= kPageSize || length % kPageSize == 0; +} + +/// \note Inline: AutoModel reads it while assembling the BackendContext, +/// before any engine exists. inline flm::backend::BackendTraits rai_traits() { flm::backend::BackendTraits traits; - traits.needs_npu_xclbin = false; // no xclbin manager is built - traits.supports_preemption = false; // rejected before the factory runs - traits.max_context_length = kRaiContextLimit; // likewise + traits.needs_npu_xclbin = false; + traits.supports_preemption = false; + traits.max_context_length = kRaiContextLimit; return traits; } -/// \brief a factory building the rai backend -/// \return a factory suitable for BackendRegistry::register_backend flm::backend::BackendFactory rai_factory(); } // namespace flm::gemma4 diff --git a/src/include/models/gemma4/rai/aie_next/gemma4_rai_constants.hpp b/src/include/models/gemma4/rai/aie_next/gemma4_rai_constants.hpp deleted file mode 100644 index d1c324b21..000000000 --- a/src/include/models/gemma4/rai/aie_next/gemma4_rai_constants.hpp +++ /dev/null @@ -1,75 +0,0 @@ -#pragma once - -#include -#include - -namespace flm::gemma4 { - -/// The PDI pair Gemma 4's ELFs were built for. Source: corelib python/ -/// conftest.py, GEMMA4_PDI. Required by ryzenai_corelib_create_stream since -/// 0.5.0 and given no default by corelib: the silicon ships the same -/// operator under several PDI tags and a model's ELFs live under exactly -/// one. A stream opened on the wrong pair silently loses every shape only -/// that pair ships, surfacing as a missing artifact rather than a wrong -/// answer -- see phi4_rai_constants.hpp's kPrefillPdi/kTokenPdi for the same -/// reasoning applied to Phi-4's own (different) pair. -inline constexpr int kPrefillPdi = 8; -inline constexpr int kTokenPdi = 17; -/// Decode PDI for a context longer than one attention window. Prefill stays -/// on kPrefillPdi; the paged kernels ship on this overlay and nowhere else. -inline constexpr int kPagedTokenPdi = 33; - -/// Shared by both shipped rows (E2B and E4B) -- read from the GGUF's own -/// `tokenizer.ggml.tokens` array length in a real load, but the two models -/// this port supports both ship this exact vocabulary. -inline constexpr std::int64_t kVocabularySize = 262144; - -/// The group size DynamicDispatch's ELF set at PDI pair (8, 17) ships at. -/// A fact about which kernels exist, not about either model -- see -/// Gemma4Row's doc comment in gemma4_rai_gguf.hpp for how this and the two -/// group fields below it were established. -inline constexpr std::uint32_t kGroupSize = 32; - -/// Longest context flat_mha ships, and the page a longer context is cut into. -inline constexpr std::int64_t kPageSize = 4096; -inline constexpr std::int64_t kMaxSequenceLength = 4096; -inline constexpr std::int64_t kMaxDecodeWindow = 4095; -/// Longest paged context: 16384, the validated ceiling. corelib upstream -/// validated paged Gemma 4 to 16k (aie4-models.md, "Past 4096 tokens"), well -/// short of the model's own 131072. A GGUF stating a shorter -/// `gemma4.context_length` caps it further. Whole pages. -inline constexpr std::int64_t kMaxPagedContext = 16384; -/// Retries for a non-finite paged prefill chunk past page 0. The chunk -/// rewrites only its own page, so a repeat is the same work. -inline constexpr int kPagedChunkAttempts = 3; - -/// \return true when `length` is 1..4096, or a multiple of kPageSize up to -/// kMaxPagedContext -inline bool ContextFits(std::uint32_t length) { - if (length < 1 || length > static_cast(kMaxPagedContext)) return false; - if (length <= static_cast(kMaxSequenceLength)) return true; - return length % static_cast(kPageSize) == 0; -} - -/// Applied to the final logits on every Gemma 4 size: tanh(x / cap) * cap. -inline constexpr float kLogitSoftcap = 30.0f; - -/// rms epsilon, read from `gemma4.attention.layer_norm_rms_epsilon` on both -/// E2B and E4B (9.999999974752427e-07, i.e. 1e-6). NOT Phi-4's 1e-5 -- -/// see phi4_rai_constants.hpp's kRmsEpsilon. Kept here only as the value a -/// fixture or a sanity check may compare against; the real loader reads the -/// metadata key rather than trusting this constant. -inline constexpr float kRmsEpsilon = 1.0e-6f; - -/// See phi4_rai_constants.hpp: 0 lets corelib choose the per-create packer -/// threads (8 at 0.25.0), on top of the concurrent creates below. -inline constexpr std::uint32_t kRequantizeThreads = 0; -inline constexpr std::size_t kWeightCreateConcurrency = 8; - -/// \brief bumped whenever the packed bytes change for an unchanged GGUF, -/// corelib version and descriptor; part of the weight-cache layout -/// \note 1: ple's gate and proj are packed from the transpose of the GGUF -/// mapping. -inline constexpr std::uint32_t kPackedLayoutRevision = 1; - -} // namespace flm::gemma4 diff --git a/src/include/models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp b/src/include/models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp deleted file mode 100644 index 5a1ef0f45..000000000 --- a/src/include/models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp +++ /dev/null @@ -1,283 +0,0 @@ -#pragma once - -#include "models/gemma4/rai/aie_next/gemma4_rai_constants.hpp" - -#include -#include -#include -#include -#include -#include -#include - -#include - -namespace flm::rai { class GgufFile; } - -namespace flm::gemma4 { - -/// \brief one row of the table: a shape to match, and what the file cannot say -/// -/// The shape fields (`layers` through `global_head_dim`) are a CLAIM about a -/// model, matched against a real GGUF's own metadata and tensor shapes by -/// `RequireRowFor`. The four fields below them -- `vocab`, `group`, -/// `head_group`, `ple_group` -- are the answer no GGUF key gives: they are -/// facts about which ELFs DynamicDispatch ships at PDI pair (8, 17), read off -/// the ELF set itself, not guessed from the model's shape. A row existing is -/// not a claim that the model runs; it is a claim that *these particular -/// group sizes* are the ones its ELFs were built at. -/// -/// Everything that differs between E2B and E4B beyond what is listed here -- -/// FFN width pattern, the shared-cache pair, which layers are full-attention -/// -- is read from the file by later tasks (see gemma4_rai_gguf.cpp for FFN -/// width and SWA-layer derivation). A row does not carry it, and that a row -/// does not need to is the evidence the design is shape-driven rather than -/// E2B-shaped: nothing here bakes in one model's numbers as if they were -/// universal. -struct Gemma4Row { - std::string_view name; - std::int64_t layers; - std::int64_t hidden; - std::int64_t q_heads; - std::int64_t kv_heads; - std::int64_t head_dim; // sliding-window layers - std::int64_t global_head_dim; // full-attention layers - std::int64_t vocab; - std::uint32_t group; - std::uint32_t head_group; - std::uint32_t ple_group; -}; - -/// \brief the rows that run -- currently E2B and E4B only -/// -/// The other three Gemma 4 sizes (1B, 12B, 27B) are deliberately absent, not -/// merely un-added: -/// - they drop KV heads on their full-attention layers specifically, so -/// one scalar `kv_heads` field cannot describe them -- a row would need -/// a per-layer-kind KV head count, which is a different (larger) shape -/// this table does not model; -/// - they have no per-layer embeddings at all, so the whole PLE tensor -/// family (`per_layer_token_embd`, `per_layer_model_proj`, the per-layer -/// PLE trio) this port depends on does not exist in their GGUFs. -/// Adding a row for one of them would not be "one more entry" -- it would be -/// a different struct. -std::span Gemma4Rows(); - -/// \brief find the row matching a file's shape -/// \throws std::runtime_error naming every shape field, when none matches -/// -/// Deliberately a lookup, not a formula: an unrecognized shape is refused -/// here, at load time, with a message that names the exact numbers and tells -/// the reader what to do next. Guessing a plausible-looking row for an -/// unknown shape would trade that clear, immediate rejection for a -/// dispatch failure far away in DynamicDispatch, over a shape whose group -/// sizes were never actually read off an ELF set -- a much harder failure -/// to trace back to "nobody added a row for this." -const Gemma4Row& RequireRowFor(std::int64_t layers, std::int64_t hidden, - std::int64_t q_heads, std::int64_t kv_heads, - std::int64_t head_dim, - std::int64_t global_head_dim); - -/// \brief the MLP width of every layer, read off blk.N.ffn_gate.weight -/// -/// \note NOT from `gemma4.feed_forward_length`. That metadata key is a -/// 35-element ARRAY on E2B and a plain SCALAR 10240 on E4B -- its -/// *type* moves between conversions of one model family. A reader -/// that learned "scalar" truncates twenty of E2B's layers (E2B is -/// 6144 for layers 0-14, then steps to 12288 from layer 15 on); one -/// that learned "array" fails to parse E4B at all. `config.json` is -/// no better: its `intermediate_size: 6144` is true of E2B's layers -/// 0-14 and wrong about 15-34 -- exactly the class of error a paper -/// table reproduces faithfully. -/// -/// The gate projection's own output shape is the only source that is -/// right about both files. In the FILE's own dims -- fastest-varying -/// first, so input width first -- `blk.N.ffn_gate.weight` is -/// `[1536, 6144]` / `[1536, 12288]` on E2B and `[2560, 10240]` on -/// E4B. `GgufFile::ShapeOf` REVERSES those, so what this function -/// reads is `ShapeOf[0]`, the OUTPUT width: `[6144, 1536]`, -/// `[12288, 1536]`, `[10240, 2560]`. Read `ShapeOf[1]` instead and -/// every layer of both models reports `hidden` -- a uniform, -/// plausible, wrong width that no shape check anywhere rejects. -/// Confirmed against the real files; see -/// test_gemma4_real_gguf.cpp, which asserts exactly these numbers -/// against gemma-4-E{2,4}B-it-Q8_0.gguf. Do not "simplify" this -/// back onto the metadata key; it is wrong on both models, in -/// opposite directions. -/// \throws std::runtime_error naming the tensor, if `blk.N.ffn_gate.weight` -/// is not 2-D. -std::vector DeriveFfnWidths(const flm::rai::GgufFile& file, - std::int64_t layers); - -/// \brief which layers use sliding-window attention, per layer -/// -/// \note Read off `blk.N.attn_q.weight`'s own OUTPUT width -- `ShapeOf[0]`, -/// since `ShapeOf` is the reverse of the file's dims -- never off a -/// derived period. Sliding layers project to `q_heads * head_dim`; full- -/// attention layers project to `q_heads * global_head_dim`. At the -/// group sizes both shipped models use that is 2048 against 4096 -- -/// the SAME two widths on E2B (hidden 1536) and E4B (hidden 2560), -/// which is exactly why width is the reliable discriminator and a -/// period is not: E2B's full-attention layers recur every 5th layer -/// (4, 9, 14, ... 34) and E4B's every 6th (5, 11, 17, ... 41). A -/// pattern generated from one model's period fits the other model's -/// weights nowhere -- it would produce a model that runs and is -/// silently wrong, computing attention with the wrong geometry on -/// most layers. Reading the width per layer, off the tensor that -/// actually carries it, needs no period at all. -/// \throws std::runtime_error naming the tensor, if `blk.N.attn_q.weight` is -/// not 2-D, or if its width matches neither the sliding nor the -/// full-attention geometry (both expected widths are named too -- -/// this is a corrupt or unrecognized file, not a guess to round -/// toward the nearer one). -std::vector DeriveLayerIsSwa(const flm::rai::GgufFile& file, - std::int64_t layers, std::int64_t q_heads, - std::int64_t head_dim, - std::int64_t global_head_dim); - -/// \brief for each layer, the index of the layer owning the KV cache it reads -/// -/// \param layer_is_swa which layers slide, from DeriveLayerIsSwa -/// \param kv_layers how many layers OWN a cache -- NOT -/// `gemma4.attention.shared_kv_layers` taken raw. That metadata key -/// counts the layers that SHARE a cache, the complement of what this -/// parameter means: `owning = block_count - shared_kv_layers` (20/18 -/// shared gives 15/24 owning, on E2B/E4B -- see -/// verified-gguf-facts.md's "correction that matters most"). Passing -/// the raw key here would search the wrong prefix for "the last -/// own-cache layer of each kind": on E4B it would yield last-full=17 -/// and last-sliding=16 instead of 23/22, and every shared layer would -/// bind the wrong cache -- a model that runs and is silently wrong. -/// Computing that subtraction belongs to whoever assembles the config -/// from the raw metadata, not to this function; this parameter takes -/// the count already corrected. -/// -/// \note A layer that shares reads THE LAST OWN-CACHE LAYER OF ITS OWN KIND -- -/// llama.cpp's rule (llama-model.cpp:2659-2667 through map_layer_ids in -/// llama-kv-cache.cpp:270). That gives 13 sliding / 14 full on E2B and -/// 22 / 23 on E4B. Neither pair follows from the other, and neither -/// follows from a layer count, so a tabulated answer would be invisibly -/// wrong on one model even though it looked right on the other. -/// \note A layer below kv_layers owns its cache and maps to itself. -/// \throws std::runtime_error if kv_layers is outside 1..layer_is_swa.size(), -/// or if a shared layer's kind (sliding or full-attention) has no -/// owning layer at all -- naming the layer and the kind. -std::vector DeriveCacheOwners(const std::vector& layer_is_swa, - std::int64_t kv_layers); - -/// \brief everything about one loaded Gemma 4 GGUF, assembled from the three -/// derivations above plus the metadata keys that need no derivation -/// -/// Every field either differs between E2B and E4B and is therefore read or -/// derived from the file (`layers`, `hidden`, `kv_heads`, `layer_*`, -/// `kv_layers`, ...), or is a fact about the ELF set no GGUF key states and -/// comes off the matched `Gemma4Row` instead (`vocab`, `group`, -/// `head_group`, `ple_group`). -struct Gemma4Config { - std::string_view row_name; // "E2B" or "E4B" - std::int64_t layers; - std::int64_t hidden; - std::int64_t q_heads; - std::int64_t kv_heads; - std::int64_t head_dim; // sliding, 256 - std::int64_t global_head_dim; // full, 512 - std::vector layer_intermediate; // DeriveFfnWidths - std::vector layer_is_swa; // DeriveLayerIsSwa - std::vector layer_cache_owner; // DeriveCacheOwners - std::int64_t kv_layers; // OWNING count -- see DeriveCacheOwners - std::int64_t sliding_window; // 512 on both - std::int64_t vocab; // from the row - std::int64_t ple_dim; - double rope_theta; // full-attention layers, 1e6 - double rope_theta_swa; // sliding layers, 1e4 - float eps; - float logit_softcap; // 30.0 - std::int32_t eos_token_id; - std::int32_t bos_token_id; - bool add_bos; - std::uint32_t group, head_group, ple_group; // from the row - std::int64_t context_length = kMaxPagedContext; // gemma4.context_length - - /// \brief the head size a layer's attention runs at -- `head_dim` if it - /// slides, `global_head_dim` if it is full-attention - std::int64_t LayerHeadDim(std::int64_t layer) const; -}; - -/// \brief assemble a Gemma4Config by reading `file`'s metadata and tensor -/// shapes, running the three derivations above, and matching the -/// result against `Gemma4Rows()` -/// -/// \note `gemma4.attention.key_length` is the FULL-attention head size and -/// `gemma4.attention.key_length_swa` is the SLIDING one -- the GGUF's -/// names run OPPOSITE to this file's (`head_dim` is sliding, -/// `global_head_dim` is full). Reading them the obvious way round -/// swaps 256 and 512 through the whole model, silently. This is the -/// one place they are read; nowhere else needs to know. -/// -/// \note `gemma4.attention.shared_kv_layers` counts the layers that SHARE a -/// cache, not the layers that OWN one -- `kv_layers` here is -/// `gemma4.block_count - gemma4.attention.shared_kv_layers`, never the -/// raw key. See DeriveCacheOwners's doc comment and -/// verified-gguf-facts.md's "correction that matters most": the raw -/// key is 20 on E2B and 18 on E4B; the owning counts this subtraction -/// produces are 15 and 24. Passing the raw key straight through would -/// search the wrong prefix for "the last own-cache layer of each -/// kind" and bind every shared layer to the wrong cache -- a model -/// that loads, runs, and is silently wrong. -/// -/// \throws std::runtime_error if `general.architecture` is not "gemma4", if -/// the file's tensor count is not `6 + layers * 17` (names both the -/// actual and the expected count), if no `Gemma4Row` matches the -/// file's shape, or if any derivation above throws. -Gemma4Config ConfigFromMetadata(const flm::rai::GgufFile& file); - -/// \brief a loaded Gemma 4 GGUF plus the config derived from it, with a -/// cross-check against the rest of the model package's JSON files -class Gemma4GgufPackage final { -public: - static std::shared_ptr Open(const std::filesystem::path& gguf_path); - ~Gemma4GgufPackage(); - - const Gemma4Config& Config() const noexcept; - const flm::rai::GgufFile& File() const noexcept; - const std::filesystem::path& Path() const noexcept; - - /// \brief cross-check the file against the package's JSON - /// \note Runs before any device object exists, so a mismatched package - /// fails at load with an exact diagnostic rather than mid-generation. - /// \note The real `config.json` is a MULTIMODAL config: `architectures`, - /// `model_type` and `eos_token_id` are top-level, but every field - /// describing the text model (`num_hidden_layers`, `vocab_size`, - /// `intermediate_size`, ...) is nested under `"text_config"`. This - /// reads those fields from `text_config`, falling back to the top - /// level only for a hypothetical flat config -- a check written - /// against a flat config rejects every real package, which is - /// exactly the failure mode this contract exists to prevent, - /// inverted. See verified-gguf-facts.md. - /// \note `num_hidden_layers` is the only one of the cross-checked fields - /// that is FATAL if missing: it is the only one capable of - /// catching an E2B/E4B swap at all, since `vocab_size` (262144) - /// and `eos_token_id` (106) are identical between the two shipped - /// rows. Those two are checked when present and tolerated when - /// absent. - /// \note Deliberately does NOT compare `config.json`'s `intermediate_size` - /// against anything: it is 6144 for E2B, which is true of layers - /// 0-14 and wrong about 15-34 (which step to 12288). Comparing it - /// to the derived widths would reject a correct package -- see - /// DeriveFfnWidths's doc comment for the same trap in the GGUF's - /// own metadata. - void ValidateGemma4Contract(const nlohmann::json& config, - const nlohmann::json& tokenizer, - const nlohmann::json& tokenizer_config) const; - -private: - struct Impl; - explicit Gemma4GgufPackage(std::unique_ptr impl); - std::unique_ptr impl_; -}; - -/// \brief the GGUF's eos first, then every other id config.json's top-level -/// `eos_token_id` array lists -std::vector ConfigEosIds(const nlohmann::json& config, std::int32_t gguf_eos); - -} // namespace flm::gemma4 diff --git a/src/include/models/gemma4/rai/aie_next/gemma4_rai_host.hpp b/src/include/models/gemma4/rai/aie_next/gemma4_rai_host.hpp deleted file mode 100644 index 8235d9d0c..000000000 --- a/src/include/models/gemma4/rai/aie_next/gemma4_rai_host.hpp +++ /dev/null @@ -1,160 +0,0 @@ -#pragma once - -#include -#include -#include -#include - -namespace flm::gemma4 { - -struct RopeTables { - std::vector cos; // [positions * (head_dim / 2)] - std::vector sin; - std::int64_t head_dim; - std::int64_t positions; -}; - -/// \brief the rotary tables for one geometry -/// \param frequency_factors `rope_freqs.weight`, one entry per rotary PAIR, -/// or empty for a regime that has none -/// \note Gemma 4 has TWO thetas -- 1e6 on full-attention layers and 1e4 on -/// sliding ones -- so this is called twice and the results are not -/// interchangeable. RoPE frequencies apply from position 0 upward, so a -/// wrong theta is wrong at position 1: the 4096 context cap does not -/// rescue it. -/// -/// \note THE FACTORS DIVIDE, AND THEY BELONG TO THE FULL-ATTENTION REGIME -/// ALONE. `config.json`'s `partial_rotary_factor: 0.25` is NOT a shorter -/// rotary -- `rope_dim == head_size` on both geometries, confirmed by -/// `gemma4.rope.dimension_count` (512) and `..._swa` (256) on both -/// shipped rows. It is a per-pair FREQUENCY SCALING carried in -/// `rope_freqs.weight`, which is why that tensor is [256] on both -/// models: one entry per pair of a 512-wide head. llama.cpp passes it as -/// `freq_factors` on full-attention layers only, where it DIVIDES -/// (`theta /= freq_factors[i/2]`), and the reference driver writes -/// `freq = freq / factors` (`gemma4_driver.py` `_rope_tables`). -/// Sliding layers pass none. -/// -/// \note WHICH DIRECTION THE DIVISION RUNS IS NOT VISIBLE IN ANY MAGNITUDE -/// ASSERTION -- both tables are bounded by 1 and both look like rotary -/// tables -- so it is pinned by the llama.cpp/driver comparison and, in -/// this suite, by the position-doubling identity in -/// TestRopeFrequencyFactorsDivideTheFrequency. -/// -/// \throws std::runtime_error naming both lengths, if `frequency_factors` is -/// non-empty and is not exactly `head_dim / 2` long -RopeTables MakeRopeTables(std::int64_t head_dim, double theta, - std::int64_t positions, - std::span frequency_factors = {}); - -/// \brief scale * x / sqrt(mean(x^2) + epsilon) -void RmsNorm(std::span x, std::span gamma, float epsilon, - std::span out); - -/// \brief decode one Q8_0 row into floats -/// \note Q8_0 is blocks of 32 int8 codes with one fp16 scale each. -void DecodeQ8Row(std::span row, std::span out); - -/// \brief round-to-nearest-even f32 -> bf16, as the kernels expect -void FloatsToBf16(std::span values, std::span out); - -/// \brief gather one token's per-layer embedding slice for every layer -/// \param per_layer_token_embd the mapped Q8_0 tensor, GGUF dims -/// `[layers * ple_dim, vocab]` -/// \param token the token id -/// \param out [layers * ple_dim] floats -/// \note per_layer_token_embd's dims are (8960, 262144) on E2B -- 2.35 -/// BILLION parameters in one tensor, which is why a model named E2B has -/// a 5.05 GB Q8_0 file, larger than Q8_0 Qwen3-8B. It stays MAPPED and -/// is gathered per token; materializing it is not an option at any -/// point. -/// \note THE LAYOUT IS `vocab` ROWS OF `layers * ple_dim`, not the transpose. -/// GGUF dims run fastest-varying first, so dims[0] -- the PLE axis -- -/// is the contiguous one and a token's whole slice is ONE ROW of 280 -/// Q8_0 blocks (E2B) or 336 (E4B). This is what the reference driver -/// does (`gemma4_driver.py`'s `per_layer_embedding()` / -/// `LazyEmbedding.rows`). -/// -/// Read the other way round -- as `layers * ple_dim` rows of `vocab`, -/// which is what this function did until Task R1 -- the TOTAL BYTE -/// COUNT IS IDENTICAL, so nothing rejects it and the result is 8960 -/// finite, plausible, wrong floats. There is no shape check anywhere -/// that can catch this; only comparing values against the reference -/// over a real file does, which test_gemma4_real_gguf.cpp now does. -/// \throws std::runtime_error if `layers * ple_dim` is not a multiple of 32 -/// (the Q8_0 block size along the contiguous axis), if `token` is -/// outside the vocabulary, if `out` is not `layers * ple_dim` wide, -/// or if the mapped tensor is shorter than `vocab` such rows. -void GatherPerLayerEmbedding(std::span per_layer_token_embd, - std::int64_t layers, std::int64_t ple_dim, - std::int64_t vocab, std::int64_t token, - std::span out); - -/// \brief the per-layer embedding input planes for a prompt, layer-major -/// -/// \param per_layer_token_embd the mapped Q8_0 `per_layer_token_embd.weight` -/// \param model_proj the mapped BF16 `per_layer_model_proj.weight`, which is -/// `layers * ple_dim` rows of `hidden` in memory -/// \param proj_norm_gamma `per_layer_proj_norm.weight`, `ple_dim` wide -/// \param scaled_embedding `[rows, hidden]`, the token embedding rows ALREADY -/// multiplied by `sqrt(hidden)` -/// \param ids the prompt, `rows` of them -/// \return `layers` planes, each `rows * ple_dim` floats -/// -/// THE ONE THING A FORWARD PASS COMPUTES ON THE HOST, and deliberately so, for -/// two independent reasons the reference driver gives (`_ple_inputs`). -/// `per_layer_model_proj` is `hidden -> layers * ple_dim` and is UNREACHABLE at -/// this model's (8, 17) PDI pair on both sizes -- E2B's 1536 -> 8960 ships at -/// group 64 only and entirely off-pair, E4B's 2560 -> 10752 at no pair and no -/// group at all -- so there is no ELF for a stream to dispatch. And llama.cpp -/// puts this same matmul on the CPU regardless: `project_per_layer_inputs` -/// says "this matrix multiplication will be performed in the input layer (i.e. -/// on the CPU)". This is not a placeholder for an NPU path. -/// -/// IT IS A GEMM, AND THE PREFILL's LARGEST HOST COST: `hidden * layers * -/// ple_dim` multiply-adds per row, 27.5 million on E4B. Keep it blocked, SIMD -/// and threaded: as one scalar dot product per (row, column) it measured ~16 ms -/// a row on the bring-up machine -- ~35 s of a 2,159-token TTFT, and ~16 ms of every -/// decode step. It ACCUMULATES IN DOUBLE on purpose: FP32 would halve the time -/// but is itself ~1e-6 over hidden = 2560, the tolerance the real-file -/// reference test holds these planes to. -/// -/// FOUR SCALE FACTORS, AND NONE OF THEM CAN MAKE THE OUTPUT NON-FINITE, so a -/// missing one is a wrong magnitude and fluent text: -/// -/// - `sqrt(ple_dim)` on the gathered per-layer embedding; -/// - `1/sqrt(hidden)` on the projection; -/// - `1/sqrt(2)` on the sum of the two; -/// - `sqrt(hidden)` on the token embedding, which belongs to the CALLER -/// (`scaled_embedding` is already scaled) because llama.cpp scales `inpL` -/// once, immediately after the gather. Applying it here as well multiplies -/// the projection by 39 and still produces finite, fluent-looking output. -/// -/// THE RMSNorm DOES NOT DIVIDE `1/sqrt(hidden)` BACK OUT, which is why it is -/// not redundant: `rms_norm(c*x) = x / sqrt(mean(x^2) + eps/c^2)`, so a uniform -/// scale RESCALES EPSILON. On this model eps is 1e-6 against a pre-norm -/// `mean(proj^2)` of 3.3e-3..7.4e-2, so the normalised rows come out at RMS -/// 0.999848 rather than 1; dropping the factor takes that to 0.999998. A small -/// effect, not an absent one. -/// -/// LAYER-MAJOR, so one layer's `[rows, ple_dim]` slice is contiguous and is a -/// single write into that layer's own device tensor. -/// -/// \throws std::runtime_error naming the offending length, if any span is not -/// the size the other arguments imply, or if a token id is outside the -/// vocabulary -std::vector> BuildPerLayerInputs( - std::span per_layer_token_embd, - std::span model_proj, - std::span proj_norm_gamma, - std::span scaled_embedding, std::span ids, - std::int64_t layers, std::int64_t ple_dim, std::int64_t hidden, - std::int64_t vocab, float epsilon); - -/// \brief tanh(x / cap) * cap, in place -/// \note Gemma 4 softcaps its logits at 30.0 on every size. Omitting this does -/// not crash -- it quietly changes the distribution, which is exactly the -/// failure the llama.cpp agreement test exists to catch. -void SoftcapLogits(std::span logits, float cap); - -} // namespace flm::gemma4 diff --git a/src/include/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.hpp b/src/include/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.hpp deleted file mode 100644 index eefdc72af..000000000 --- a/src/include/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.hpp +++ /dev/null @@ -1,153 +0,0 @@ -#pragma once - -#include "models/gemma4/rai/aie_next/gemma4_rai_constants.hpp" -#include "models/gemma4/rai/aie_next/gemma4_rai_gguf.hpp" -#include "rai/corelib_api.hpp" - -#include -#include -#include -#include -#include - -namespace flm::gemma4 { - -/// \brief the padded row extent every op needs for one bucket, one geometry -struct Gemma4RowExtents { - std::int64_t q_rows; - std::int64_t kv_rows; - std::int64_t output_rows; - std::int64_t ssmlp_rows; - /// \note Not queried through corelib -- see gemma4_rai_shape_plan.cpp's - /// Build() for why `ryzenai_corelib_rmsnorm_bf16_pad_rows` cannot - /// be trusted for this even though it exists. - std::int64_t rmsnorm_rows; - /// \note Not queried through corelib either -- there is no - /// `ryzenai_corelib_ple_bf16_pad_rows` in the pinned 0.5.0 header. - /// See Build()'s own comment. - std::int64_t ple_rows; - std::int64_t flat_mha_rows; -}; - -/// \brief asks corelib how it wants Gemma 4's operators padded, for BOTH -/// attention geometries at once, and holds the FOUR attention -/// descriptors a Gemma 4 layer loop can need -/// -/// TWO AXES, AND THEY ARE NOT THE SAME KIND OF THING. -/// -/// The first is GEOMETRY. Gemma 4 runs two in one stack: sliding-window layers -/// at head_dim 256, full-attention layers at head_dim 512. They resolve -/// DIFFERENT ELFs (`ryzenai_corelib_flat_mha_bf16_desc::head_size` is part of -/// what selects the kernel) and they genuinely have different SHAPES, so a -/// single row-extent table cannot answer for both -- `ForRows(rows, swa)` -/// keeps one per geometry. -/// -/// The second is CACHE ROLE: whether a layer owns the KV cache it attends -/// over or reads one an earlier layer wrote (`Gemma4Config::layer_cache_owner`). -/// `kv_shared` selects the "_kvshare" artifacts, which are handed no K operand -/// at all and merge nothing into the cache (corelib.h:1772-1779). This is a -/// 2x2, not a second copy of the geometry split. -/// -/// BUT THE ROW EXTENTS DO NOT SPLIT ON THE SECOND AXIS. The header states it -/// outright, as the very reason `flat_mha_bf16_elf_name` exists: the -/// descriptor's "window, kv_shared and scale do not change the shapes of -/// anything -- they change WHICH KERNEL RUNS, by changing the name DD -/// resolves" (corelib.h:1810-1812). So `flat_mha_bf16_pad_rows` cannot depend -/// on `kv_shared`, the owning and sharing families share one row grid per -/// geometry, and `sliding_rows_`/`full_rows_` are correct as two tables rather -/// than four. -/// -/// WHY ALL FOUR DESCRIPTORS ARE MATERIALIZED ANYWAY. An earlier version of -/// this plan held two templates with `kv_shared = 0` and left the layer loop -/// to override the field per layer before dispatching. The reference driver -/// keys four descriptors on `(sliding, shares-a-cache)` instead -/// (`ryzenai-corelib/python/gemma4_driver.py:1526-1559`) and says why: "An -/// owning layer given a kvshare descriptor never writes its own cache; a -/// sharing layer given an owning one overwrites somebody else's. Neither -/// errors." A mutated shared template makes that a one-missed-assignment bug -/// with no failure mode; four immutable descriptors and a selecting accessor -/// make it unrepresentable. -class Gemma4ShapePlan final { -public: - /// \brief interrogate corelib for the padded extent of every live row - /// count, for both geometries, plus one ssmlp query per distinct - /// FFN width the model has - /// \param api the corelib binding - /// \param stream the stream the plan will be dispatched on -- corelib - /// 0.5.0 takes the stream on every padding helper, because the PDI - /// pair a stream was opened with selects the ELF set, and a shape - /// exists under one pair and not another - /// \param config the model's shape, read off its GGUF (Task C5) - /// \param context the engine's context. Above kPageSize the plan is paged: - /// attention descriptors carry that window and a kPageSize page. - static Gemma4ShapePlan Build( - const std::shared_ptr& api, - ryzenai_corelib_stream_ptr stream, const Gemma4Config& config, - std::int64_t context = kMaxSequenceLength); - - bool paged() const noexcept { return paged_; } - - /// \param live_rows the live (unpadded) row count a caller wants extents - /// for, 1..4096 - /// \param swa true for a sliding-window layer, false for full attention - /// \note Held PER GEOMETRY because head_dim 256 and 512 resolve DIFFERENT - /// ELFs. One table cannot answer for both. - const Gemma4RowExtents& ForRows(std::size_t live_rows, bool swa) const; - - /// \param swa true for a sliding-window layer, false for full attention - /// \param shared true for a layer that READS another layer's KV cache - /// \note SELECT, never mutate. The returned descriptor is one of four the - /// plan built once; assigning `kv_shared` on a copy held across a - /// layer loop is the bug this accessor exists to remove -- see the - /// class comment. - /// \note THERE IS DELIBERATELY NO ONE-ARGUMENT OVERLOAD. An earlier draft - /// of this fix kept `attention_desc(swa)` meaning "the owning one", - /// which reintroduces the failure in a quieter form: a layer loop - /// that writes `attention_desc(layer_is_swa[i])` and forgets the - /// cache role compiles, runs, and hands every sharing layer an - /// owning descriptor. Requiring the second argument turns that into - /// a compile error. Callers that genuinely want the owning one say - /// `attention_desc(swa, /*shared=*/false)`. - const ryzenai_corelib_flat_mha_bf16_desc& attention_desc( - bool swa, bool shared) const noexcept; - - /// \note Valid only when paged() is true. Same (geometry, cache role) - /// index as attention_desc. - const ryzenai_corelib_paged_mha_bf16_desc& paged_attention_desc( - bool swa, bool shared) const noexcept; - - /// \brief the padded row extent lm_head's matmul runs at -- the leading - /// extent of the `[rows, vocab]` logits tensor a caller allocates - /// \note ONE, on every shipped Gemma 4, and NOT a simplification. lm_head - /// is `hidden -> vocab` and ships at M == 1 alone at this model's - /// (8, 17) PDI pair: "there is no batched form to offer even if a - /// caller wanted every position's logits" - /// (`ryzenai-corelib/python/gemma4_driver.py:2449-2453`). It is - /// still asked of corelib rather than hardcoded, because the answer - /// is corelib's to give. See Build(). - std::int64_t lm_head_rows() const noexcept; - -private: - /// \brief index into `attention_descs_` for one (geometry, cache role) - /// \note A free-standing constexpr rather than an inline expression at - /// each use, so the accessor and Build() cannot drift into - /// disagreeing about the ordering. - static constexpr std::size_t DescIndex(bool swa, bool shared) noexcept { - return (swa ? 0u : 2u) + (shared ? 1u : 0u); - } - - std::vector sliding_rows_; - std::vector full_rows_; - /// \brief lm_head's padded row extent -- see lm_head_rows() - /// \note NOT part of Gemma4RowExtents: those are per row bucket and per - /// geometry, and lm_head is neither. It has one shape. - std::int64_t lm_head_rows_{0}; - /// \brief the four (geometry, cache role) descriptors, by DescIndex - /// \note Materialized once in Build() and never written again. See the - /// class comment for why four and not two. - std::array attention_descs_{}; - std::array paged_descs_{}; - bool paged_{false}; -}; - -} // namespace flm::gemma4 diff --git a/src/include/models/qwen3/rai/aie_next/qwen3_rai.hpp b/src/include/models/qwen3/rai/aie_next/qwen3_rai.hpp deleted file mode 100644 index 192ca418e..000000000 --- a/src/include/models/qwen3/rai/aie_next/qwen3_rai.hpp +++ /dev/null @@ -1,48 +0,0 @@ -#pragma once - -#include "causal_lm.hpp" -#include "lm_config.hpp" -#include "models/qwen3/rai/aie_next/qwen3_rai_gguf.hpp" -#include "rai/corelib_runtime.hpp" - -#include -#include - -namespace flm::qwen3 { - -/// \brief Qwen3 0.6B/1.7B/4B/8B on ryzenai-corelib, from a Q8_0 GGUF -/// \note corelib has no chunked prefill -- a multi-row pass must start at -/// position 0 -- so the engine remembers the tokens behind its KV cache, -/// and a prefill that continues a conversation re-runs them from the -/// front together with the new ones. -class qwen3_rai final : public causal_lm { -public: - qwen3_rai(LM_Config config, - std::shared_ptr package, - std::shared_ptr runtime, - std::uint32_t max_length = 4096); - ~qwen3_rai() override; - - buffer forward(int id) override; - buffer prefill(std::vector& ids, void* payload = nullptr) override; - void set_context_length(int length) override; - /// \brief unsupported; this engine's weights come from its GGUF package - /// \note Present only because causal_lm.hpp is a frozen ABI. See the - /// definition and AutoModel/model_backend.hpp. - /// \throws std::runtime_error always - void load_weights(Q4NX&) override; - void update_max_length(std::uint32_t max_length) override; - void clear_context() override; - buffer get_k_cache(int layer, int index) override; - buffer get_v_cache(int layer, int index) override; - int get_current_context_length() override; - int checkpoint() override; - int restore() override; - bool poisoned() const noexcept; - -private: - struct Impl; - std::unique_ptr impl_; -}; - -} // namespace flm::qwen3 diff --git a/src/include/models/qwen3/rai/aie_next/qwen3_rai_backend.hpp b/src/include/models/qwen3/rai/aie_next/qwen3_rai_backend.hpp index 9f32e935e..839734dc5 100644 --- a/src/include/models/qwen3/rai/aie_next/qwen3_rai_backend.hpp +++ b/src/include/models/qwen3/rai/aie_next/qwen3_rai_backend.hpp @@ -12,9 +12,6 @@ namespace flm::qwen3 { /// \brief the largest context this backend can hold inline constexpr std::uint32_t kRaiContextLimit = 4096; -/// \brief the largest number of tokens corelib will decode -inline constexpr std::uint32_t kRaiDecodeLimit = 4095; - /// \brief what the frontend must know before building this backend /// \note Inline so the registry can reject a request before the engine exists. inline flm::backend::BackendTraits rai_traits() { diff --git a/src/include/models/qwen3/rai/aie_next/qwen3_rai_config.hpp b/src/include/models/qwen3/rai/aie_next/qwen3_rai_config.hpp deleted file mode 100644 index 84bf047dd..000000000 --- a/src/include/models/qwen3/rai/aie_next/qwen3_rai_config.hpp +++ /dev/null @@ -1,107 +0,0 @@ -/// \file qwen3_rai_config.hpp -/// \brief The Qwen3 sizes the rai engine runs, and the kernel grid it runs on -/// \note One engine serves four sizes, so the architecture is a row of a table -/// rather than a block of constants. The GGUF's own metadata selects the -/// row and is cross-checked against it; the row supplies what the file -/// cannot state -- the group sizes, which are properties of which ELFs -/// ship, not of the model. -#pragma once - -#include -#include -#include -#include -#include -#include - -namespace flm::qwen3 { - -struct Qwen3Config { - std::string_view size; - std::int64_t layers; - std::int64_t hidden; - std::int64_t q_heads; - std::int64_t kv_heads; - std::int64_t head_dim; - std::int64_t intermediate; - std::int64_t vocab; - double rope_theta; - float epsilon; - /// \brief the requantized group of every projection and the MLP - std::uint32_t group; - /// \brief the group lm_head packs at, which is a different ELF family - std::uint32_t head_group; - - std::int64_t q_dim() const noexcept { return q_heads * head_dim; } - std::int64_t kv_dim() const noexcept { return kv_heads * head_dim; } -}; - -/// \brief the four sizes ryzenai-corelib's qwen3 driver runs on the NPU -/// \note Transcribed from corelib's QWEN table (python/model_driver/ -/// qwen3_driver.py), which is where the claim that they run is kept. -/// 0.6B is the only row whose lm_head packs at group 32: 1024 -> 151936 -/// ships a group-32 kernel only, while everything else in that model -/// ships group 64 only. -inline constexpr std::array kQwen3Rows{{ - {"0.6B", 28, 1024, 16, 8, 128, 3072, 151936, 1.0e6, 1.0e-6f, 64, 32}, - {"1.7B", 28, 2048, 16, 8, 128, 6144, 151936, 1.0e6, 1.0e-6f, 64, 64}, - {"4B", 36, 2560, 32, 8, 128, 9728, 151936, 1.0e6, 1.0e-6f, 64, 64}, - {"8B", 36, 4096, 32, 8, 128, 12288, 151936, 1.0e6, 1.0e-6f, 64, 64}, -}}; - -/// \brief the arch fields a GGUF states, as read from its metadata -struct Qwen3Shape { - std::int64_t layers, hidden, q_heads, kv_heads, head_dim, intermediate; - double rope_theta; - double epsilon; -}; - -/// \brief the row a file's shape selects -/// \throws std::runtime_error when no row matches, or when the matching row -/// disagrees with the file about epsilon or rope theta (Qwen3-4B-2507 -/// is shape-identical to Qwen3-4B but ropes at 5e6, not 1e6) -const Qwen3Config& SelectQwen3Row(const Qwen3Shape& shape); - -/// \brief rows buffers and the attention ELF are built for; a power of two -inline constexpr std::int64_t kMaxSequenceLength = 4096; -/// \brief the last position a decode may write -inline constexpr std::int64_t kMaxDecodeWindow = 4095; - -/// \brief the row counts every op these models dispatch ships a kernel for -/// \note matmul, ssmlp and flat_mha share this grid at every covered width, -/// which is what lets one padded row count serve a whole pass. The shape -/// plan asks corelib and refuses to load if any op disagrees. -inline constexpr std::array kExecutionRows{ - 1, 64, 128, 256, 512, 1024, 2048, 3072, 4096}; - -/// \brief QK-Norm's width, and every M a k=128 RMSNorm kernel ships at -/// \note Read off dyn_bins.dll's RMSNorm kernels (token PDI and prefill PDI): -/// 8/16/32 on the token PDI, 512 and up on prefill, and 33..511 a hole -/// no PDI covers. corelib's rmsnorm pad helper is an identity and cannot -/// report this, so the shape plan checks rows x heads against it. -inline constexpr std::int64_t kQkNormWidth = 128; -inline constexpr std::array kQkNormRows{ - 8, 16, 32, 512, 1024, 2048, 4096, 8192, 16384, 24576, 32768, 49152, - 65536, 98304, 131072}; - -/// \brief the PDI pair every Qwen3 shape ships under (corelib's PDI_PAIR) -inline constexpr int kPrefillPdi = 1; -inline constexpr int kTokenPdi = 16; - -/// \brief intra-packer thread hint; 0 lets corelib choose -inline constexpr std::uint32_t kRequantizeThreads = 0; - -/// \brief how many requantizing creates run at once -/// \note Serial, as corelib's qwen3 driver loads. corelib records an open, -/// unattributed defect on both requantizing entry points in which -/// concurrent creates correlate with all-zero device output; serial is -/// the configuration it has not been observed in. The weight cache makes -/// this a first-load cost only. -inline constexpr std::size_t kWeightCreateConcurrency = 1; - -/// \brief Qwen3's two end ids: <|im_end|> ends a turn, <|endoftext|> a document -/// \note The GGUF and tokenizer_config.json each state only 151645; the pair is -/// Qwen's generation_config.json. Both are checked against tokenizer.json. -inline constexpr std::array kEosIds{151645, 151643}; - -} // namespace flm::qwen3 diff --git a/src/include/models/qwen3/rai/aie_next/qwen3_rai_gguf.hpp b/src/include/models/qwen3/rai/aie_next/qwen3_rai_gguf.hpp deleted file mode 100644 index 365f8fb4d..000000000 --- a/src/include/models/qwen3/rai/aie_next/qwen3_rai_gguf.hpp +++ /dev/null @@ -1,60 +0,0 @@ -/// \file qwen3_rai_gguf.hpp -/// \brief A Qwen3 Q8_0 GGUF, validated and resolved before any device work -#pragma once - -#include "models/qwen3/rai/aie_next/qwen3_rai_config.hpp" -#include "rai/gguf_file.hpp" - -#include -#include -#include - -#include - -namespace flm::qwen3 { - -/// \brief one decoder block's tensors, as spans over the mapping -/// \note Nothing is fused in a Qwen3 conversion: Q, K and V are three tensors -/// and gate and up two, so each view is a whole tensor. -struct Qwen3LayerTensors { - rai::GgufFloatTensor attn_norm, ffn_norm, q_norm, k_norm; - rai::GgufTensor q, k, v, o, gate, up, down; -}; - -class Qwen3GgufPackage final { -public: - /// \brief map the file, select its size and shape-check every tensor - /// \throws std::runtime_error naming the field at the first mismatch - static std::shared_ptr Open(const std::filesystem::path& gguf_path); - ~Qwen3GgufPackage(); - - const Qwen3Config& Config() const noexcept; - const std::filesystem::path& Path() const; - const rai::GgufTensor& Embedding() const noexcept; - /// \brief the tensor lm_head packs from - /// \note A property of the file, not the family: Qwen3-8B ships its own - /// output.weight, the smaller three tie lm_head to token_embd. The - /// two are the same shape, so reading the wrong one raises nothing. - const rai::GgufTensor& LmHead() const noexcept; - bool TiedLmHead() const noexcept; - const rai::GgufFloatTensor& OutputNorm() const noexcept; - const Qwen3LayerTensors& Layer(std::size_t layer) const; - - /// \brief cross-check the file against the package's JSON files - /// \note config.json is checked only for the keys it states: FastFlowLM's - /// Qwen3-4B config omits tie_word_embeddings. - void ValidateContract(const nlohmann::json& config, - const nlohmann::json& tokenizer, - const nlohmann::json& tokenizer_config) const; - -private: - Qwen3GgufPackage() = default; - - std::shared_ptr file_; - const Qwen3Config* config_{}; - rai::GgufTensor embedding_{}, lm_head_{}; - rai::GgufFloatTensor output_norm_{}; - std::vector layers_; -}; - -} // namespace flm::qwen3 diff --git a/src/include/models/qwen3/rai/aie_next/qwen3_rai_shape_plan.hpp b/src/include/models/qwen3/rai/aie_next/qwen3_rai_shape_plan.hpp deleted file mode 100644 index 2f5a2dcef..000000000 --- a/src/include/models/qwen3/rai/aie_next/qwen3_rai_shape_plan.hpp +++ /dev/null @@ -1,36 +0,0 @@ -/// \file qwen3_rai_shape_plan.hpp -/// \brief The padded row count every op of one pass runs at -#pragma once - -#include "models/qwen3/rai/aie_next/qwen3_rai_config.hpp" -#include "rai/corelib_api.hpp" - -#include -#include -#include - -namespace flm::qwen3 { - -class Qwen3ShapePlan final { -public: - /// \brief interrogate corelib at every execution bucket - /// \throws std::runtime_error when any op pads a bucket to something else, - /// when a helper moves K or N, or when rows x heads lands in the - /// k=128 RMSNorm hole - /// \note Qwen3 dispatches matmul, QK-Norm, flat_mha and ssmlp over the same - /// buffers, so they must agree on one row count; asking corelib, on - /// the stream the plan is for, turns a grid assumption into a load - /// failure rather than a wrong answer mid-generation. - static Qwen3ShapePlan Build(const std::shared_ptr& api, - ryzenai_corelib_stream_ptr stream, - const Qwen3Config& config); - - /// \brief the bucket a pass of `live_rows` runs at - std::int64_t RowsFor(std::size_t live_rows) const; - const ryzenai_corelib_flat_mha_bf16_desc& attention_desc() const noexcept { return attention_; } - -private: - ryzenai_corelib_flat_mha_bf16_desc attention_{}; -}; - -} // namespace flm::qwen3 diff --git a/src/include/models/qwen35/rai/aie_next/qwen35_rai.hpp b/src/include/models/qwen35/rai/aie_next/qwen35_rai.hpp deleted file mode 100644 index 243991ce5..000000000 --- a/src/include/models/qwen35/rai/aie_next/qwen35_rai.hpp +++ /dev/null @@ -1,43 +0,0 @@ -/// \file qwen35_rai.hpp -/// \brief Qwen3.5 4B and 9B text on ryzenai-corelib, from a Q8_0 GGUF -#pragma once - -#include "causal_lm.hpp" -#include "lm_config.hpp" -#include "models/qwen35/rai/aie_next/qwen35_rai_gguf.hpp" -#include "rai/corelib_runtime.hpp" - -#include -#include - -namespace flm::qwen35 { - -/// \brief dense Qwen3.5 on the NPU through ryzenai-corelib -/// \note Vision and MTP are not this engine. Prefill steps the decode chain, -/// which is the reference the batched prefill path is checked against. -class qwen35_rai final : public causal_lm { -public: - qwen35_rai(LM_Config config, std::shared_ptr package, - std::shared_ptr runtime, - std::uint32_t max_length = kMaxSequenceLength); - ~qwen35_rai() override; - - buffer forward(int id) override; - buffer prefill(std::vector& ids, void* payload = nullptr) override; - void set_context_length(int length) override; - void load_weights(Q4NX&) override; - void update_max_length(std::uint32_t max_length) override; - void clear_context() override; - buffer get_k_cache(int layer, int index) override; - buffer get_v_cache(int layer, int index) override; - int get_current_context_length() override; - int checkpoint() override; - int restore() override; - bool poisoned() const noexcept; - -private: - struct Impl; - std::unique_ptr impl_; -}; - -} // namespace flm::qwen35 diff --git a/src/include/models/qwen35/rai/aie_next/qwen35_rai_backend.hpp b/src/include/models/qwen35/rai/aie_next/qwen35_rai_backend.hpp index e7f57e6c9..45475bcaf 100644 --- a/src/include/models/qwen35/rai/aie_next/qwen35_rai_backend.hpp +++ b/src/include/models/qwen35/rai/aie_next/qwen35_rai_backend.hpp @@ -10,7 +10,6 @@ namespace flm::qwen35 { inline constexpr std::uint32_t kRaiContextLimit = 4096; -inline constexpr std::uint32_t kRaiDecodeLimit = 4095; inline flm::backend::BackendTraits rai_traits() { flm::backend::BackendTraits traits; diff --git a/src/include/models/qwen35/rai/aie_next/qwen35_rai_gguf.hpp b/src/include/models/qwen35/rai/aie_next/qwen35_rai_gguf.hpp deleted file mode 100644 index 6dda442d3..000000000 --- a/src/include/models/qwen35/rai/aie_next/qwen35_rai_gguf.hpp +++ /dev/null @@ -1,103 +0,0 @@ -/// \file qwen35_rai_gguf.hpp -/// \brief A Qwen3.5 Q8_0 GGUF, with the linear-attention V heads back in HF order -#pragma once - -#include "rai/gguf_file.hpp" - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace flm::qwen35 { - -inline constexpr int kPrefillPdi = 9; -inline constexpr int kTokenPdi = 19; -inline constexpr int kGroup = 64; -inline constexpr int kMaxSequenceLength = 4096; -inline constexpr int kSoftplusOperand = 64; -inline constexpr int kRequantizeThreads = 1; -inline constexpr std::array kEosIds{248046, 248044}; - -struct Qwen35Config { - int layers = 0; - int hidden = 0; - int intermediate = 0; - int q_heads = 0; - int kv_heads = 0; - int head_dim = 0; - int rope_dim = 0; - double rope_theta = 0; - double eps = 0; - int full_every = 0; - int conv_width = 0; - int lin_k_heads = 0; - int lin_v_heads = 0; - int lin_k_dim = 0; - int lin_v_dim = 0; - int vocab = 0; - bool tied_lm_head = false; - - int q_dim() const { return q_heads * head_dim; } - int kv_dim() const { return kv_heads * head_dim; } - int lin_qk_dim() const { return lin_k_heads * lin_k_dim; } - int lin_value_dim() const { return lin_v_heads * lin_v_dim; } - int conv_channels() const { return 2 * lin_qk_dim() + lin_value_dim(); } - bool is_full(int layer) const { return (layer + 1) % full_every == 0; } -}; - -/// \brief group-64 ONNX MatMulNBits components for one biased projection -struct OnnxQ8 { - std::vector qweight; - std::vector scales; - std::vector qzeros; -}; - -class Qwen35Gguf { -public: - static std::shared_ptr Open(const std::filesystem::path& path); - - const std::filesystem::path& Path() const { return path_; } - const Qwen35Config& Config() const { return config_; } - const std::string& LmHeadName() const { return lm_head_; } - const rai::GgufTensor& Embedding() const { return embedding_; } - - /// \brief raw Q8_0 bytes for an [n, k] matrix, k fastest - std::span Blocks(std::string_view base, std::int64_t k, std::int64_t n) const; - std::span F32(std::string_view name, std::span shape) const; - - std::vector QkvBlocks(int layer) const; - std::vector ZBlocks(int layer) const; - std::vector GateBlocks(int layer, std::string_view which) const; - std::vector OutBlocks(int layer) const; - std::vector PerHead(int layer, std::string_view name) const; - /// \brief [channels, width] float32, channels already in HF order - std::vector ConvTaps(int layer) const; - -private: - Qwen35Gguf(std::shared_ptr file, std::filesystem::path path, - Qwen35Config config, std::string lm_head, rai::GgufTensor embedding, - std::vector order); - - std::vector TakeRows(std::span raw, std::int64_t n, std::int64_t k, - std::span rows) const; - std::vector ChannelOrder() const; - - std::shared_ptr file_; - std::filesystem::path path_; - Qwen35Config config_; - std::string lm_head_; - rai::GgufTensor embedding_; - std::vector order_; -}; - -/// \brief refit Q8_0 [n, k] into the ONNX layout from_onnx packs, group 64 -/// \note Byte-identical to gguf_requantized where both apply. Used only for the -/// alpha gate, which needs a bias the GGUF requantizer has no slot for. -OnnxQ8 Q8ToOnnx(std::span blocks, int k, int n, int group); - -} // namespace flm::qwen35 diff --git a/src/include/models/qwen36/rai/aie_next/qwen36_rai.hpp b/src/include/models/qwen36/rai/aie_next/qwen36_rai.hpp deleted file mode 100644 index 585f58d5c..000000000 --- a/src/include/models/qwen36/rai/aie_next/qwen36_rai.hpp +++ /dev/null @@ -1,46 +0,0 @@ -/// \file qwen36_rai.hpp -/// \brief Qwen3.6-35B-A3B text on ryzenai-corelib, from a Q8_0 GGUF -#pragma once - -#include "causal_lm.hpp" -#include "rai/corelib_runtime.hpp" - -#include -#include -#include - -namespace flm::qwen36 { - -inline constexpr std::uint32_t kContextLimit = 262144; - -/// \brief Qwen3.6 text through corelib's model API -/// \note Vision is not this engine. A prompt is one forward(); corelib runs -/// its first 4096 rows from position 0 as one pass and steps the rest. -/// The KV cache is sized once, at the capacity `max_length` rounds up -/// to in whole 4096-row pages. -class qwen36_rai final : public causal_lm { -public: - qwen36_rai(std::filesystem::path gguf, - std::shared_ptr runtime, - std::uint32_t max_length = 4096); - ~qwen36_rai() override; - - buffer forward(int id) override; - buffer prefill(std::vector& ids, void* payload = nullptr) override; - void set_context_length(int length) override; - void load_weights(Q4NX&) override; - void update_max_length(std::uint32_t max_length) override; - void clear_context() override; - buffer get_k_cache(int layer, int index) override; - buffer get_v_cache(int layer, int index) override; - int get_current_context_length() override; - int checkpoint() override; - int restore() override; - bool poisoned() const noexcept; - -private: - struct Impl; - std::unique_ptr impl_; -}; - -} // namespace flm::qwen36 diff --git a/src/include/models/qwen36/rai/aie_next/qwen36_rai_backend.hpp b/src/include/models/qwen36/rai/aie_next/qwen36_rai_backend.hpp index 5665ae526..e8b11a67e 100644 --- a/src/include/models/qwen36/rai/aie_next/qwen36_rai_backend.hpp +++ b/src/include/models/qwen36/rai/aie_next/qwen36_rai_backend.hpp @@ -10,7 +10,6 @@ namespace flm::qwen36 { inline constexpr std::uint32_t kRaiContextLimit = 262144; -inline constexpr std::uint32_t kRaiDecodeLimit = kRaiContextLimit - 1; inline flm::backend::BackendTraits rai_traits() { flm::backend::BackendTraits traits; diff --git a/src/include/rai/corelib_api.hpp b/src/include/rai/corelib_api.hpp index 2c57f4ab9..4e89cd655 100644 --- a/src/include/rai/corelib_api.hpp +++ b/src/include/rai/corelib_api.hpp @@ -31,32 +31,18 @@ X(create_host_view, ryzenai_corelib_create_host_view) \ X(tensor_write, ryzenai_corelib_tensor_write) \ X(tensor_read, ryzenai_corelib_tensor_read) \ - X(tensor_get_byte_size, ryzenai_corelib_tensor_get_byte_size) \ - X(tensor_get_data_type, ryzenai_corelib_tensor_get_data_type) \ - X(stream_get_kernels_root, ryzenai_corelib_stream_get_kernels_root) \ - X(stream_set_kernels_root, ryzenai_corelib_stream_set_kernels_root) \ X(stream_enable_kernel_bundle, ryzenai_corelib_stream_enable_kernel_bundle) \ - X(weights_slice, ryzenai_corelib_weights_slice) \ X(matmul_enum_kernels, ryzenai_corelib_matmul_bf16_enum_kernels) \ X(matmul_weights_create, ryzenai_corelib_matmul_bf16_weights_create) \ X(matmul, ryzenai_corelib_matmul_bf16) \ - X(matmul_act, ryzenai_corelib_matmul_act_bf16) \ X(matmul_act_enum_kernels, ryzenai_corelib_matmul_act_bf16_enum_kernels) \ X(ssmlp_enum_kernels, ryzenai_corelib_ssmlp_bf16_enum_kernels) \ X(ssmlp_weights_create, ryzenai_corelib_ssmlp_bf16_weights_create) \ X(weights_copy_data, ryzenai_corelib_weights_copy_data) \ X(ssmlp, ryzenai_corelib_ssmlp_bf16) \ - X(rmsnorm_weights_create, ryzenai_corelib_rmsnorm_bf16_weights_create) \ - X(rmsnorm, ryzenai_corelib_rmsnorm_bf16) \ - X(ple_weights_create, ryzenai_corelib_ple_bf16_weights_create) \ - X(ple, ryzenai_corelib_ple_bf16) \ X(flat_mha_enum_kernels, ryzenai_corelib_flat_mha_bf16_enum_kernels) \ X(flat_mha, ryzenai_corelib_flat_mha_bf16) \ X(paged_mha_enum_kernels, ryzenai_corelib_paged_mha_bf16_enum_kernels) \ - X(paged_mha, ryzenai_corelib_paged_mha_bf16) \ - X(linear_attention, ryzenai_corelib_linear_attention_bf16) \ - X(dwconv_weights_create, ryzenai_corelib_dwconv_bf16_weights_create) \ - X(dwconv, ryzenai_corelib_dwconv_bf16) \ X(model_info_get, ryzenai_corelib_model_info_get) \ X(model_forward, ryzenai_corelib_model_forward) \ X(model_reset, ryzenai_corelib_model_reset) \ diff --git a/src/include/rai/corelib_model.hpp b/src/include/rai/corelib_model.hpp new file mode 100644 index 000000000..b45466233 --- /dev/null +++ b/src/include/rai/corelib_model.hpp @@ -0,0 +1,87 @@ +/// \file corelib_model.hpp +/// \brief A causal_lm over a model corelib runs whole (ryzenai_corelib_model_*) +#pragma once + +#include "causal_lm.hpp" +#include "rai/corelib_runtime.hpp" + +#include +#include +#include +#include + +namespace flm::rai { + +/// \brief where a multi-row pass may start when the model is past position 0 +enum class Continuation { + /// corelib takes rows > 1 at any position + kAnywhere, + /// only at position 0: a continuation re-runs every earlier token with it + kFromStart, + /// at 0 or, on a paged cache, on a page boundary: a continuation re-runs + /// the earlier tokens of its own page with it + kFromPage, +}; + +/// \brief how the model goes back to an earlier position +enum class Rewind { + /// attention only: each pass states its position and overwrites what lies past it + kSeek, + /// recurrent state in model slots: model_state_save/restore at the + /// checkpoint, reset and replay anywhere else + kSnapshot, + /// recurrent state outside any slot: reset and replay + kReplay, +}; + +/// \brief what tells one corelib model apart from another +struct ModelSpec { + /// \brief corelib's name for the model, e.g. "qwen3" + std::string corelib_name; + /// \brief how messages name the model, e.g. "Qwen3" + std::string display_name; + /// \brief the longest max_length this engine accepts + std::uint32_t context_limit; + Continuation continuation; + Rewind rewind; +}; + +/// \brief rows in one page of a paged KV cache +inline constexpr std::uint32_t kPageRows = 4096; + +/// \brief one corelib model, its packed weights, and the tokens it holds +/// \note Text only. corelib packs the weights from the file at construction, +/// which is the slow part of a load; the KV cache is sized once, at the +/// capacity corelib rounds `max_length` up to. +class CorelibModel final : public causal_lm { +public: + /// \param weights the file corelib packs the weights from + /// \param max_length positions this engine may run: 1..spec.context_limit + CorelibModel(ModelSpec spec, std::filesystem::path weights, + std::shared_ptr runtime, + std::uint32_t max_length); + ~CorelibModel() override; + + buffer forward(int id) override; + buffer prefill(std::vector& ids, void* payload = nullptr) override; + void set_context_length(int length) override; + void load_weights(Q4NX&) override; + void update_max_length(std::uint32_t max_length) override; + void clear_context() override; + buffer get_k_cache(int layer, int index) override; + buffer get_v_cache(int layer, int index) override; + int get_current_context_length() override; + int checkpoint() override; + int restore() override; + + /// \brief whether a corelib call failed mid-sequence + /// \note Nothing clears it; every other entry point throws while it is + /// set, and the only way out is unload and reload. + bool poisoned() const noexcept; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace flm::rai diff --git a/src/include/rai/corelib_model_backend.hpp b/src/include/rai/corelib_model_backend.hpp new file mode 100644 index 000000000..026367720 --- /dev/null +++ b/src/include/rai/corelib_model_backend.hpp @@ -0,0 +1,36 @@ +/// \file corelib_model_backend.hpp +/// \brief The rai backend for a model corelib runs whole, one spec per family +/// \note Compiled only when FLM_ENABLE_RAI is on. +#pragma once + +#include "AutoModel/model_backend.hpp" +#include "rai/corelib_model.hpp" + +#include +#include +#include +#include + +namespace flm::rai { + +struct BackendSpec { + ModelSpec model; + /// \brief whether a requested context is one the model runs; empty means + /// 1..model.context_limit + std::function context_fits; + /// \brief the contexts allowed, as a refusal states them, e.g. "1..4096" + std::string context_rule; + /// \brief the ids generation stops on; empty reads them from the package: + /// tokenizer_config.json's eos_token, then config.json's eos_token_id + std::vector eos_ids; + /// \brief report tokenizer_config.json's bos_token as the BOS id, for a + /// package that names the token but not its id (Gemma 4) + bool read_bos = false; +}; + +/// \brief a factory building the rai backend `spec` describes +/// \note Preemption, the context and the package's token ids are checked +/// before the runtime is acquired or the weights are opened. +flm::backend::BackendFactory Factory(BackendSpec spec); + +} // namespace flm::rai diff --git a/src/include/rai/corelib_object.hpp b/src/include/rai/corelib_object.hpp index 5f7a72da3..1e2eb65ef 100644 --- a/src/include/rai/corelib_object.hpp +++ b/src/include/rai/corelib_object.hpp @@ -14,9 +14,6 @@ struct TensorWindowTag {}; struct HostViewTag {}; struct MatMulWeightsTag {}; struct SsMlpWeightsTag {}; -struct RmsNormWeightsTag {}; -struct PleWeightsTag {}; -struct DwConvWeightsTag {}; struct ModelTag {}; struct ModelWeightsTag {}; struct ModelSnapshotTag {}; @@ -69,9 +66,6 @@ using UniqueTensorWindow = UniqueObject; using UniqueHostView = UniqueObject; using UniqueMatMulWeights = UniqueObject; using UniqueSsMlpWeights = UniqueObject; -using UniqueRmsNormWeights = UniqueObject; -using UniquePleWeights = UniqueObject; -using UniqueDwConvWeights = UniqueObject; using UniqueModel = UniqueObject; using UniqueModelWeights = UniqueObject; using UniqueModelSnapshot = UniqueObject; diff --git a/src/include/rai/gguf_file.hpp b/src/include/rai/gguf_file.hpp deleted file mode 100644 index df7343b65..000000000 --- a/src/include/rai/gguf_file.hpp +++ /dev/null @@ -1,101 +0,0 @@ -/// \file gguf_file.hpp -/// \brief A memory-mapped GGUF v3 file, for the rai engines that read one -/// \note Only what a corelib engine needs: the tensor directory, as spans over -/// the mapping, and the scalar metadata a caller asks to keep. Arrays are -/// walked and skipped (only their element count is kept), which is what -/// keeps opening a file with a 150k-entry vocabulary cheap. -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace flm::rai { - -inline constexpr std::uint32_t kGgmlTypeF32 = 0; -inline constexpr std::uint32_t kGgmlTypeQ8_0 = 8; -inline constexpr std::uint32_t kGgmlTypeBF16 = 30; - -/// \brief a tensor as the file stores it -/// \note `shape` is outermost-first -- the reverse of GGUF's own dimension -/// order -- so a projection from k to n reads [n, k]. -struct GgufTensor { - std::string_view name; - std::span bytes; - std::vector shape; - std::uint32_t ggml_type; -}; - -/// \brief an F32 tensor, viewed as floats in place -struct GgufFloatTensor { - std::string_view name; - std::span values; - std::vector shape; -}; - -/// \brief a BF16 tensor, viewed in place as raw 16-bit patterns -struct GgufBf16Tensor { - std::string_view name; - std::span values; - std::vector shape; - std::uint32_t ggml_type; -}; - -class GgufFile final { -public: - /// \brief which metadata keys to keep; everything else is parsed and dropped - using RetainKey = std::function; - - /// \brief map and index a GGUF v3 file - /// \param retain which keys to keep; empty keeps every key - /// \throws std::runtime_error naming the field when the file is malformed: - /// bad magic or version, ranges outside the file, misaligned or - /// overlapping tensors, duplicate names - static std::shared_ptr Open(const std::filesystem::path& path, - const RetainKey& retain = {}); - ~GgufFile(); - - const std::filesystem::path& Path() const; - bool HasTensor(std::string_view name) const; - bool HasMetadata(std::string_view key) const; - - /// \brief a Q8_0 tensor of exactly this shape - GgufTensor RequireQ8(std::string_view name, - std::span shape) const; - /// \brief an F32 tensor of exactly this shape - GgufFloatTensor RequireF32(std::string_view name, - std::span shape) const; - /// \brief a BF16 tensor of exactly this shape - GgufBf16Tensor RequireBf16(std::string_view name, - std::span shape) const; - /// \brief whatever the file holds under this name, unchecked - GgufTensor Tensor(std::string_view name) const; - std::size_t TensorCount() const noexcept; - - std::uint64_t Unsigned(std::string_view key) const; - double Number(std::string_view key) const; - bool Boolean(std::string_view key) const; - std::string String(std::string_view key) const; - /// \brief the element count of an array-valued key - std::uint64_t ArrayCount(std::string_view key) const; - /// \brief the elements of an integer array, widened - /// \throws when the elements are not integers or one is negative - std::vector UnsignedArray(std::string_view key) const; - -private: - struct Impl; - explicit GgufFile(std::unique_ptr impl); - std::unique_ptr impl_; -}; - -/// \brief "field: actual X, expected Y", the shape every rai load error takes -[[noreturn]] void FailField(std::string_view field, std::string actual, - std::string expected); - -} // namespace flm::rai diff --git a/src/include/rai/host_ops.hpp b/src/include/rai/host_ops.hpp index 92d2de560..2328040cd 100644 --- a/src/include/rai/host_ops.hpp +++ b/src/include/rai/host_ops.hpp @@ -3,23 +3,12 @@ /// \note Plain CPU and no corelib types, so it is testable without hardware. #pragma once -#include "rai/gguf_file.hpp" - #include #include #include namespace flm::rai { -/// \brief dequantize the rows of a [vocab, width] Q8_0 table for these ids -/// \return ids.size() * width floats, row-major -std::vector DecodeQ8Rows(const GgufTensor& table, std::span ids); - -/// \brief scale * x / sqrt(mean(x^2) + epsilon), row by row -void RmsNorm(std::span input, std::span scale, - std::int64_t rows, std::int64_t width, float epsilon, - std::span output); - /// \brief round-to-nearest-even F32 -> BF16, NaN kept quiet std::vector F32ToBf16(std::span values); diff --git a/src/include/rai/weight_cache.hpp b/src/include/rai/weight_cache.hpp deleted file mode 100644 index 02065dbca..000000000 --- a/src/include/rai/weight_cache.hpp +++ /dev/null @@ -1,80 +0,0 @@ -/// \file weight_cache.hpp -/// \brief On-disk cache of a model's packed corelib weights -/// \note Requantizing every weight from Q8_0 is effectively the whole of a rai -/// model load. corelib hands the packed bytes back -/// (ryzenai_corelib_weights_copy_data) and maps them in again later -/// (..._weights_create_from_file), so the refit is paid once per GGUF. -/// One data file plus a JSON index written last and renamed into place: -/// a data file without a matching index is never used. -#pragma once - -#include -#include -#include -#include -#include - -namespace flm::rai { - -/// \brief where one packed weight lives inside the cache file -struct CachedWeightSpan { - std::uint64_t offset{}; - std::uint64_t size{}; -}; - -/// \brief the identity a cache is only valid for -/// \note The GGUF is identified by size and write time, not content: hashing a -/// multi-gigabyte file costs more than the packing the cache avoids. -/// corelib rejects a slice whose length is not exactly what its -/// descriptor packs to, so a stale-but-plausible cache is caught there. -/// `layout` names everything else that shapes the packed bytes -- every -/// group size, every width -- so changing any of them is a miss. -struct WeightCacheKey { - std::uint64_t gguf_size{}; - std::int64_t gguf_write_time{}; - std::uint32_t corelib_major{}, corelib_minor{}, corelib_patch{}; - std::string layout; - std::uint64_t weight_count{}; - - bool operator==(const WeightCacheKey& other) const = default; -}; - -WeightCacheKey MakeWeightCacheKey(const std::filesystem::path& gguf_path, - std::uint32_t corelib_major, - std::uint32_t corelib_minor, - std::uint32_t corelib_patch, - std::string layout, - std::uint64_t weight_count); - -class WeightCache final { -public: - /// \brief the cache for one GGUF, or nullopt when caching is disabled - /// \param gguf_path the model file; the cache is named after it - /// \note FLM_RAI_WEIGHT_CACHE unset puts the cache beside the GGUF; a path - /// redirects it; "0", "off" or "false" disables it. Named after the - /// GGUF, so several models can share one redirected directory. - static std::optional ForGguf(const std::filesystem::path& gguf_path); - - WeightCache(std::filesystem::path directory, std::string stem); - - std::filesystem::path DataPath() const; - std::filesystem::path IndexPath() const; - - /// \brief the spans, when an index is present, current and in bounds - /// \note Never throws: a damaged cache is a miss, not a failed load. - std::optional> ReadIndex(const WeightCacheKey& expected) const; - - /// \brief write the index for a data file already renamed into place - /// \return false when it could not be written - bool WriteIndex(const WeightCacheKey& key, const std::vector& spans) const; - - /// \brief delete this cache and any temporaries an interrupted write left - /// \return the bytes reclaimed; never throws - std::uint64_t Remove() const; - -private: - std::filesystem::path directory_; - std::string stem_; -}; - -} // namespace flm::rai diff --git a/src/test/gemma4_rai/CMakeLists.txt b/src/test/gemma4_rai/CMakeLists.txt index 2c224060c..3dd2bfb1a 100644 --- a/src/test/gemma4_rai/CMakeLists.txt +++ b/src/test/gemma4_rai/CMakeLists.txt @@ -12,34 +12,10 @@ endif() include_directories("${FLM_SOURCE_DIR}/include" "${FLM_SOURCE_DIR}/test/phi4_rai") -# Every Gemma 4 test target is declared here up front, each guarded on the -# existence of its own .cpp, so later tasks add only their source file and -# never edit this file -- two agents touching this CMakeLists at once is -# exactly the collision a previous attempt on this project produced. -# -# Order and provenance (see docs/superpowers/plans/2026-09-22-gemma4-rai.md, -# Stage A and Stage C): -# test_gemma4_host Task A2 (this task) -# test_gemma4_frontend Task A3 -# test_gemma4_config Task C1 -# test_gemma4_derivations, test_gemma4_gguf Tasks C2-C5 -# test_gemma4_shape_plan Task C6 -# test_gemma4_engine Tasks C7-C9 -# test_gemma4_backend Task C10 - -# --- test_gemma4_host (Task A2) ----------------------------------------- -# The host layer holds NO corelib types -- that is the whole point of it, -# and it is what lets this target be built and run with no NPU, no XRT, no -# boost and no RYZENAI_CORELIB_INCLUDE_DIR. It links nothing but its own -# .cpp. If this target ever starts needing the corelib header or an XRT/ -# boost include path, the host layer has been contaminated: fix the layer, -# not this file. -if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_host.cpp") - add_executable(test_gemma4_host - test_gemma4_host.cpp - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_host.cpp") - add_test(NAME test_gemma4_host COMMAND test_gemma4_host) -endif() +# The Gemma 4 test targets: +# test_gemma4_frontend the Gemma4e frontend, over a stub engine +# test_gemma4_tokenizer the shipped tokenizer stack on the real packages +# test_gemma4_registry the rai backend, registered under gemma4e # --- test_gemma4_frontend (Task A3) -------------------------------------- # Exercises Gemma4e::engine_config() through the real AutoModel/Gemma4e @@ -49,8 +25,8 @@ endif() # any.h pulls in boost/any.hpp) but not the corelib header itself, since the # frontend never includes rai/corelib_api.hpp directly. if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_frontend.cpp") - # Not REQUIRED: a configure run that only wants test_gemma4_host (or any - # other dependency-free target) must not be blocked by this target's + # Not REQUIRED: a configure run that only wants another target here + # must not be blocked by this target's # probes. Absence is reported below -- loudly, and as a ctest FAILURE -- # rather than failing the whole configure. # @@ -96,11 +72,11 @@ if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_frontend.cpp") target_link_libraries(test_gemma4_frontend PRIVATE xrt_coreutil) add_test(NAME test_gemma4_frontend COMMAND test_gemma4_frontend) else() - # LOUD AT CONFIGURE, AND RED AT CTEST -- the treatment C6 gave - # test_gemma4_shape_plan and C7 gave test_gemma4_engine, arriving here - # last. message(STATUS) alone is visible while configuring and - # invisible afterwards: the ctest run that follows says "100% tests - # passed" and nothing in it names the target that was never compiled. + # LOUD AT CONFIGURE, AND RED AT CTEST -- the convention every block + # in this file follows. message(STATUS) alone is visible while + # configuring and invisible afterwards: the ctest run that follows + # says "100% tests passed" and nothing in it names the target that + # was never compiled. # A green run that covers less than it appears to is the same species # as the vacuous test, the self-confirming fixture and the stale .exe, # and this project has now been bitten by all four. @@ -119,79 +95,6 @@ if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_frontend.cpp") endif() endif() -# --- test_gemma4_config (Task C1) ---------------------------------------- -# Reads a Gemma4Config off a GGUF fixture. gemma4_rai_gguf.cpp is Task C1's -# own file; flm::rai::GgufFile (rai/gguf_file.cpp) is Task B1's. Neither -# includes a corelib header -- phi4_rai_gguf.cpp, the file this one is -# modelled on, does not either -- so this target needs no RYZENAI_CORELIB_ -# INCLUDE_DIR and no fake_corelib. If the owning task's design ends up -# needing corelib here after all, add the find_path/link the same way -# test_gemma4_shape_plan does below. -if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_config.cpp") - add_executable(test_gemma4_config - test_gemma4_config.cpp - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp" - "${FLM_SOURCE_DIR}/common/rai/gguf_file.cpp") - target_include_directories(test_gemma4_config PRIVATE - "${CMAKE_CURRENT_LIST_DIR}" "${FLM_SOURCE_DIR}/include") - add_test(NAME test_gemma4_config COMMAND test_gemma4_config) -endif() - -# --- test_gemma4_derivations (Tasks C2-C4) ------------------------------- -# Same sources as test_gemma4_config: FFN width, SWA layer period and cache -# ownership are all derivations read off the same gemma4_rai_gguf.cpp. -if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_derivations.cpp") - add_executable(test_gemma4_derivations - test_gemma4_derivations.cpp - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp" - "${FLM_SOURCE_DIR}/common/rai/gguf_file.cpp") - target_include_directories(test_gemma4_derivations PRIVATE - "${CMAKE_CURRENT_LIST_DIR}" "${FLM_SOURCE_DIR}/include") - add_test(NAME test_gemma4_derivations COMMAND test_gemma4_derivations) -endif() - -# --- test_gemma4_gguf (Task C5) ------------------------------------------ -# The whole-package contract (Gemma4GgufPackage), same sources again. -if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_gguf.cpp") - add_executable(test_gemma4_gguf - test_gemma4_gguf.cpp - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp" - "${FLM_SOURCE_DIR}/common/rai/gguf_file.cpp") - target_include_directories(test_gemma4_gguf PRIVATE - "${CMAKE_CURRENT_LIST_DIR}" "${FLM_SOURCE_DIR}/include") - add_test(NAME test_gemma4_gguf COMMAND test_gemma4_gguf) -endif() - -# --- test_gemma4_real_gguf (Task R1) ------------------------------------- -# THE ONLY TARGET HERE THAT READS A FILE GOOGLE SHIPPED. Everything above -# reads gemma4_gguf_fixture.hpp's synthetic file, and for four signed-off -# tasks that fixture and the readers written against it confirmed each other -# while both were transposed relative to every real conversion. This target -# is the answer to that; see test_gemma4_real_gguf.cpp's header comment. -# -# It is OPT-IN on two environment variables naming the real GGUFs, because -# they are 4.8 GB and 7.7 GB and will not exist on every machine. The skip -# convention is the tree's existing one for external artifacts -- -# src/test/phi4_rai's test_real_corelib: the binary returns 77 and ctest -# prints "Skipped" on its own line, which is loud, unlike a target that -# silently is not built. -# -# This block IS an edit to a file whose header says later tasks should only -# add their own .cpp. That instruction exists so two agents working in -# parallel do not collide here; there was no block for a real-file target -# because until this task no real file could be opened at all. -if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_real_gguf.cpp") - add_executable(test_gemma4_real_gguf - test_gemma4_real_gguf.cpp - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_host.cpp" - "${FLM_SOURCE_DIR}/common/rai/gguf_file.cpp") - target_include_directories(test_gemma4_real_gguf PRIVATE - "${CMAKE_CURRENT_LIST_DIR}" "${FLM_SOURCE_DIR}/include") - add_test(NAME test_gemma4_real_gguf COMMAND test_gemma4_real_gguf) - set_tests_properties(test_gemma4_real_gguf PROPERTIES SKIP_RETURN_CODE 77) -endif() - # --- test_gemma4_tokenizer (Task R5) ------------------------------------- # THE BINARY THAT TAKES A STEP OFF THE HARDWARE CHECKLIST. review-R4.md §5 # settled the doubled-BOS question on the host and recommended landing it as @@ -266,7 +169,7 @@ if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_tokenizer.cpp") # third_party/tokenizers-cpp/CMakeLists.txt's own list. ntdll wsock32 ws2_32 Bcrypt iphlpapi userenv psapi) add_test(NAME test_gemma4_tokenizer COMMAND test_gemma4_tokenizer) - # Opt-in on the real packages, exactly like test_gemma4_real_gguf. + # Opt-in on the real packages: without them it exits 77, a skip. set_tests_properties(test_gemma4_tokenizer PROPERTIES SKIP_RETURN_CODE 77) else() # LOUD AT CONFIGURE, AND RED AT CTEST -- the convention every block in @@ -296,270 +199,55 @@ if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_tokenizer.cpp") endif() endif() -# --- test_gemma4_shape_plan (Task C6) ------------------------------------ -# Builds corelib descriptors (flat_mha_bf16_desc, ssmlp) per geometry, so -# this one does need the corelib header and fake_corelib, the same way -# test_phi4_shape_plan does in src/test/phi4_rai/CMakeLists.txt. -if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_shape_plan.cpp") - # Not REQUIRED -- see the comment on test_gemma4_frontend's probes above; - # same reasoning applies to every target below. - # - # HINTS, though, and they are not cosmetic. This probe had none, so a - # configure that did not pass -DRYZENAI_CORELIB_INCLUDE_DIR found nothing, - # skipped this target, and then `ctest` reported "100% tests passed ... out - # of 4" -- a green run over four targets that is indistinguishable, in a CI - # line or a scrollback, from the green run over five. A review reproduced - # exactly that. The pinned 0.5.0 header lives in a sibling worktree's - # .deps, which is where the second and third hints point. - find_path(RYZENAI_CORELIB_INCLUDE_DIR NAMES ryzenai/corelib.h - HINTS "$ENV{RYZENAI_CORELIB_INCLUDE_DIR}" - "${CMAKE_CURRENT_LIST_DIR}/../../../.deps/corelib-050/include") - if(RYZENAI_CORELIB_INCLUDE_DIR) - add_executable(test_gemma4_shape_plan - test_gemma4_shape_plan.cpp - "${CMAKE_CURRENT_LIST_DIR}/fake_corelib.cpp" - "${FLM_SOURCE_DIR}/common/rai/corelib_api.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp" - "${FLM_SOURCE_DIR}/common/rai/gguf_file.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp") - target_include_directories(test_gemma4_shape_plan PRIVATE - "${CMAKE_CURRENT_LIST_DIR}" - "${FLM_SOURCE_DIR}/include" - "${RYZENAI_CORELIB_INCLUDE_DIR}") - target_compile_definitions(test_gemma4_shape_plan PRIVATE RYZENAI_CORELIB_STATIC=1) - add_test(NAME test_gemma4_shape_plan COMMAND test_gemma4_shape_plan) - else() - # LOUD AT CONFIGURE, AND RED AT CTEST. message(STATUS) alone -- which - # is what this branch used to be -- is visible when configuring and - # invisible afterwards: the ctest run still says "100% tests passed" - # and nothing in it names the target that was never compiled. A green - # run that covers less than it appears to is the same species as the - # vacuous test, the self-confirming fixture and the stale .exe, and - # this project has been bitten by all four. - # - # So: WARNING rather than STATUS, plus a registered test that FAILS. - # The test's command succeeds and prints why; FAIL_REGULAR_EXPRESSION - # is what turns that into a ctest failure, so the explanation lands in - # the failure output instead of having to be hunted for in the - # configure log. - # - # This is deliberately NOT find_path(... REQUIRED): a configure that - # only wants the dependency-free targets (test_gemma4_host and - # friends) must still succeed, which is the convention the blocks - # above established. It only refuses to let such a run LOOK complete. - message(WARNING "gemma4_rai: test_gemma4_shape_plan will NOT be built -- " - "RYZENAI_CORELIB_INCLUDE_DIR not found (need ryzenai/corelib.h). " - "Pass -DRYZENAI_CORELIB_INCLUDE_DIR= or set it in the " - "environment. ctest will report test_gemma4_shape_plan_WAS_NOT_BUILT " - "as FAILED so this run cannot be mistaken for a complete one.") - add_test(NAME test_gemma4_shape_plan_WAS_NOT_BUILT - COMMAND "${CMAKE_COMMAND}" -E echo - "test_gemma4_shape_plan WAS NOT BUILT: RYZENAI_CORELIB_INCLUDE_DIR " - "was not found at configure time, so this ctest run does not " - "cover it. Re-configure with -DRYZENAI_CORELIB_INCLUDE_DIR=.") - set_tests_properties(test_gemma4_shape_plan_WAS_NOT_BUILT PROPERTIES - FAIL_REGULAR_EXPRESSION "WAS NOT BUILT") - endif() -endif() - -# --- test_gemma4_engine (Tasks C7-C9) ------------------------------------ -# The full engine: weights, the layer loop, teardown/poisoned(). Needs -# corelib, boost and XRT the same way test_phi4_engine does. -if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_engine.cpp") - # HINTS, AND THEY ARE NOT COSMETIC -- this probe had none, exactly like - # test_gemma4_shape_plan's did before C6 fixed it. A configure that did - # not pass -DRYZENAI_CORELIB_INCLUDE_DIR found nothing, skipped the - # target with a message(STATUS) nobody sees afterwards, and then `ctest` - # reported "100% tests passed" over a set that did not include the - # engine. A green run that covers less than it appears to is the failure - # this project keeps paying for. - find_path(RYZENAI_CORELIB_INCLUDE_DIR NAMES ryzenai/corelib.h - HINTS "$ENV{RYZENAI_CORELIB_INCLUDE_DIR}" - "${CMAKE_CURRENT_LIST_DIR}/../../../.deps/corelib-050/include") - find_path(BOOST_INCLUDE_DIR NAMES boost/program_options.hpp - HINTS "$ENV{CONDA_PREFIX}/Library/include" - "$ENV{USERPROFILE}/anaconda3/Library/include") - # The fifth hint is for a GIT WORKTREE. The fourth resolves to - # /xrt_package only when the checkout sits directly in the - # workspace; a worktree under /.worktrees/ is two levels - # deeper, so the same package needs two more "..". Without it this - # target silently vanished in every worktree, which is where this port - # is actually developed. - find_path(XRT_INCLUDE_DIR NAMES xrt/xrt_bo.h - HINTS "$ENV{XRT_INCLUDE_DIR}" - "${CMAKE_CURRENT_LIST_DIR}/../../../../xrt_package/xrt/include" - "${CMAKE_CURRENT_LIST_DIR}/../../../../../../xrt_package/xrt/include" - "C:/dev/XRT/src/runtime_src/core/include") - if(RYZENAI_CORELIB_INCLUDE_DIR AND BOOST_INCLUDE_DIR AND XRT_INCLUDE_DIR) - add_executable(test_gemma4_engine - test_gemma4_engine.cpp - "${CMAKE_CURRENT_LIST_DIR}/fake_corelib.cpp" - "${FLM_SOURCE_DIR}/common/rai/corelib_api.cpp" - "${FLM_SOURCE_DIR}/common/rai/corelib_runtime.cpp" - "${FLM_SOURCE_DIR}/common/rai/gguf_file.cpp" - "${FLM_SOURCE_DIR}/common/rai/weight_cache.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_host.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai.cpp") - target_include_directories(test_gemma4_engine PRIVATE - "${CMAKE_CURRENT_LIST_DIR}" - "${FLM_SOURCE_DIR}/include" - "${RYZENAI_CORELIB_INCLUDE_DIR}" - "${BOOST_INCLUDE_DIR}" - "${XRT_INCLUDE_DIR}") - target_compile_definitions(test_gemma4_engine PRIVATE - RYZENAI_CORELIB_STATIC=1 USEAVX2=1 DISABLE_ABI_CHECK=1 - _ENABLE_EXTENDED_ALIGNED_STORAGE WIN32_LEAN_AND_MEAN NOMINMAX - # Task C9's structural teardown gate READS THE ENGINE'S SOURCE. - # The order corelib requires -- every packed object released - # before the Stream -- is a property of `Impl`'s MEMBER - # DECLARATION ORDER, and the fault it prevents is intermittent - # and raises nothing, so the only test with teeth is one that - # reads the declarations themselves. The test FAILS, loudly, if - # this path does not open; a structural gate that skips when it - # cannot find its subject is worse than none. - "FLM_GEMMA4_RAI_SOURCE=\"${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai.cpp\"") - target_compile_options(test_gemma4_engine PRIVATE - $<$:/wd4005 /wd4244>) - target_link_directories(test_gemma4_engine PRIVATE "${XRT_INCLUDE_DIR}/../lib") - target_link_libraries(test_gemma4_engine PRIVATE xrt_coreutil) - add_test(NAME test_gemma4_engine COMMAND test_gemma4_engine) - else() - # LOUD AT CONFIGURE, AND RED AT CTEST -- the same treatment C6 gave - # test_gemma4_shape_plan, and for the same reason: message(STATUS) - # alone is invisible once the configure log has scrolled past, and - # the ctest run that follows says "100% tests passed" without naming - # the target that was never compiled. - message(WARNING "gemma4_rai: test_gemma4_engine will NOT be built -- " - "RYZENAI_CORELIB_INCLUDE_DIR=${RYZENAI_CORELIB_INCLUDE_DIR}, " - "BOOST_INCLUDE_DIR=${BOOST_INCLUDE_DIR}, XRT_INCLUDE_DIR=${XRT_INCLUDE_DIR} " - "(need ryzenai/corelib.h, boost/program_options.hpp, xrt/xrt_bo.h). " - "ctest will report test_gemma4_engine_WAS_NOT_BUILT as FAILED so " - "this run cannot be mistaken for a complete one.") - add_test(NAME test_gemma4_engine_WAS_NOT_BUILT - COMMAND "${CMAKE_COMMAND}" -E echo - "test_gemma4_engine WAS NOT BUILT: one of ryzenai/corelib.h, " - "boost/program_options.hpp or xrt/xrt_bo.h was not found at " - "configure time, so this ctest run does not cover it.") - set_tests_properties(test_gemma4_engine_WAS_NOT_BUILT PROPERTIES - FAIL_REGULAR_EXPRESSION "WAS NOT BUILT") - endif() -endif() - -# --- test_gemma4_backend (Task C10) -------------------------------------- -# The backend, its traits, and registration under family gemma4e. Same -# dependency shape as test_gemma4_engine, plus the backend's own .cpp. -if(EXISTS "${CMAKE_CURRENT_LIST_DIR}/test_gemma4_backend.cpp") - # HINTS, AND THEY ARE NOT COSMETIC -- this block was written - # speculatively, before C6 and C7 had found what a hintless probe costs. - # The corelib probe had NO hints at all, so a configure that did not pass - # -DRYZENAI_CORELIB_INCLUDE_DIR found nothing and dropped this target with - # a message(STATUS) nobody sees afterwards; and the XRT probe was missing - # the fifth hint, the one for a GIT WORKTREE (the fourth resolves to - # /xrt_package only when the checkout sits directly in the - # workspace; a worktree under /.worktrees/ is two levels - # deeper). Either way `ctest` then reports "100% tests passed" over a set - # that does not include the backend. Fixed here the same way C6 fixed - # test_gemma4_shape_plan and C7 fixed test_gemma4_engine. - find_path(RYZENAI_CORELIB_INCLUDE_DIR NAMES ryzenai/corelib.h - HINTS "$ENV{RYZENAI_CORELIB_INCLUDE_DIR}" - "${CMAKE_CURRENT_LIST_DIR}/../../../.deps/corelib-050/include") - find_path(BOOST_INCLUDE_DIR NAMES boost/program_options.hpp - HINTS "$ENV{CONDA_PREFIX}/Library/include" - "$ENV{USERPROFILE}/anaconda3/Library/include") - find_path(XRT_INCLUDE_DIR NAMES xrt/xrt_bo.h - HINTS "$ENV{XRT_INCLUDE_DIR}" - "${CMAKE_CURRENT_LIST_DIR}/../../../../xrt_package/xrt/include" - "${CMAKE_CURRENT_LIST_DIR}/../../../../../../xrt_package/xrt/include" - "C:/dev/XRT/src/runtime_src/core/include") - if(RYZENAI_CORELIB_INCLUDE_DIR AND BOOST_INCLUDE_DIR AND XRT_INCLUDE_DIR) - add_executable(test_gemma4_backend - test_gemma4_backend.cpp - "${CMAKE_CURRENT_LIST_DIR}/fake_corelib.cpp" - # The registry itself, so registration, traits lookup and - # resolve_backend_id all run their production code. The one thing - # this target does NOT link is builtin_backends.cpp: that file - # names every engine type and would drag in the ~20 prebuilt - # libraries under src/lib/. The test defines its own - # empty register_builtin_backends instead -- the same trick - # src/test/phi4_rai/test_phi4_frontend.cpp uses. - "${FLM_SOURCE_DIR}/common/AutoModel/model_backend.cpp" - "${FLM_SOURCE_DIR}/common/rai/corelib_api.cpp" - "${FLM_SOURCE_DIR}/common/rai/corelib_runtime.cpp" - "${FLM_SOURCE_DIR}/common/rai/gguf_file.cpp" - "${FLM_SOURCE_DIR}/common/rai/weight_cache.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_gguf.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_host.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_shape_plan.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai.cpp" - "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp") - target_include_directories(test_gemma4_backend PRIVATE - "${CMAKE_CURRENT_LIST_DIR}" - "${FLM_SOURCE_DIR}/include" - "${RYZENAI_CORELIB_INCLUDE_DIR}" - "${BOOST_INCLUDE_DIR}" - "${XRT_INCLUDE_DIR}") - target_compile_definitions(test_gemma4_backend PRIVATE - RYZENAI_CORELIB_STATIC=1 USEAVX2=1 DISABLE_ABI_CHECK=1 - _ENABLE_EXTENDED_ALIGNED_STORAGE WIN32_LEAN_AND_MEAN NOMINMAX - # The registration this target asserts is the one builtin_backends - # .cpp makes under `#if defined(FLM_ENABLE_RAI)`, so the guard has - # to be on here too -- otherwise the assertion is about a branch - # the shipping build does not take. - FLM_ENABLE_RAI=1 - # Turns flm::file_access's observer on. Without it ObserveOpen is - # an empty inline and "nothing was opened before the rejection" - # is vacuously true -- the exact shape of vacuous test this - # project has been bitten by. - FLM_CORELIB_TESTING=1 - # The BACKEND's teardown gate READS ITS OWN SOURCE, the same way - # C9's engine gate reads gemma4_rai.cpp. "The runtime shared_ptr - # is declared before the engine member" is what makes the engine - # release before the runtime, it is invisible at runtime (the - # fault corelib raises is intermittent and raises nothing here), - # and a later edit could reorder two member declarations without - # a single test noticing. The test FAILS, loudly, if this path - # does not open. - "FLM_GEMMA4_RAI_BACKEND_SOURCE=\"${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp\"" - # C-2's two gates read their subjects the same way. Neither - # function can be LINKED here: get_auto_model includes every - # frontend in the tree, and ensure_model_loaded is the whole - # server. What this binary can still say is that the server - # refuses an unknown model name before it resolves it, and that - # get_auto_model no longer picks a frontend from a literal tag. - # Both tests FAIL, loudly, if these paths do not open. - "FLM_REST_HANDLER_SOURCE=\"${FLM_SOURCE_DIR}/server/rest_handler.cpp\"" - "FLM_ALL_MODELS_SOURCE=\"${FLM_SOURCE_DIR}/include/AutoModel/all_models.hpp\"" - # The SHIPPED catalog, for the two tripwires D1 left: both are - # about fallback_model() becoming reachable for a family that now - # has a rai backend, and neither can be asked of a hand-written - # catalog -- the whole question is what the real model_list.json - # resolves to on aie_next. - "FLM_TEST_MODEL_LIST_PATH=\"${FLM_SOURCE_DIR}/model_list.json\"") - target_compile_options(test_gemma4_backend PRIVATE - $<$:/wd4005 /wd4244>) - target_link_directories(test_gemma4_backend PRIVATE "${XRT_INCLUDE_DIR}/../lib") - target_link_libraries(test_gemma4_backend PRIVATE xrt_coreutil) - add_test(NAME test_gemma4_backend COMMAND test_gemma4_backend) - else() - # LOUD AT CONFIGURE, AND RED AT CTEST -- this block shipped with a - # bare message(STATUS), which is visible while configuring and - # invisible afterwards: the ctest run that follows says "100% tests - # passed" and nothing in it names the target that was never compiled. - # Same treatment as every other block above, for the same reason. - message(WARNING "gemma4_rai: test_gemma4_backend will NOT be built -- " - "RYZENAI_CORELIB_INCLUDE_DIR=${RYZENAI_CORELIB_INCLUDE_DIR}, " - "BOOST_INCLUDE_DIR=${BOOST_INCLUDE_DIR}, XRT_INCLUDE_DIR=${XRT_INCLUDE_DIR} " - "(need ryzenai/corelib.h, boost/program_options.hpp, xrt/xrt_bo.h). " - "ctest will report test_gemma4_backend_WAS_NOT_BUILT as FAILED so " - "this run cannot be mistaken for a complete one.") - add_test(NAME test_gemma4_backend_WAS_NOT_BUILT - COMMAND "${CMAKE_COMMAND}" -E echo - "test_gemma4_backend WAS NOT BUILT: one of ryzenai/corelib.h, " - "boost/program_options.hpp or xrt/xrt_bo.h was not found at " - "configure time, so this ctest run does not cover it.") - set_tests_properties(test_gemma4_backend_WAS_NOT_BUILT PROPERTIES - FAIL_REGULAR_EXPRESSION "WAS NOT BUILT") - endif() +# --- test_gemma4_registry ---------------------------------------------------- +# The rai backend's traits, its registration under gemma4e, and the shipped +# catalog and server gates. The engine is corelib's own model behind phi4_rai's +# fake, so nothing here says the model is right. +find_path(RYZENAI_CORELIB_INCLUDE_DIR NAMES ryzenai/corelib.h + HINTS "$ENV{RYZENAI_CORELIB_INCLUDE_DIR}") +find_path(BOOST_INCLUDE_DIR NAMES boost/program_options.hpp + HINTS "$ENV{CONDA_PREFIX}/Library/include" + "$ENV{USERPROFILE}/anaconda3/Library/include") +find_path(XRT_INCLUDE_DIR NAMES xrt/xrt_bo.h + HINTS "$ENV{XRT_INCLUDE_DIR}" + "${CMAKE_CURRENT_LIST_DIR}/../../../../xrt_package/xrt/include" + "${CMAKE_CURRENT_LIST_DIR}/../../../../../../xrt_package/xrt/include" + "C:/dev/XRT/src/runtime_src/core/include") +if(RYZENAI_CORELIB_INCLUDE_DIR AND BOOST_INCLUDE_DIR AND XRT_INCLUDE_DIR) + add_executable(test_gemma4_registry + test_gemma4_registry.cpp + "${FLM_SOURCE_DIR}/test/phi4_rai/fake_corelib.cpp" + "${FLM_SOURCE_DIR}/common/AutoModel/model_backend.cpp" + "${FLM_SOURCE_DIR}/common/rai/corelib_api.cpp" + "${FLM_SOURCE_DIR}/common/rai/corelib_runtime.cpp" + "${FLM_SOURCE_DIR}/common/rai/host_ops.cpp" + "${FLM_SOURCE_DIR}/common/rai/corelib_model.cpp" + "${FLM_SOURCE_DIR}/common/rai/corelib_model_backend.cpp" + "${FLM_SOURCE_DIR}/common/models/gemma4/rai/aie_next/gemma4_rai_backend.cpp") + target_include_directories(test_gemma4_registry PRIVATE + "${CMAKE_CURRENT_LIST_DIR}" + "${FLM_SOURCE_DIR}/include" + "${RYZENAI_CORELIB_INCLUDE_DIR}" + "${BOOST_INCLUDE_DIR}" + "${XRT_INCLUDE_DIR}") + target_compile_definitions(test_gemma4_registry PRIVATE + RYZENAI_CORELIB_STATIC=1 USEAVX2=1 DISABLE_ABI_CHECK=1 + _ENABLE_EXTENDED_ALIGNED_STORAGE WIN32_LEAN_AND_MEAN NOMINMAX + FLM_ENABLE_RAI=1 FLM_CORELIB_TESTING=1 + "FLM_REST_HANDLER_SOURCE=\"${FLM_SOURCE_DIR}/server/rest_handler.cpp\"" + "FLM_ALL_MODELS_SOURCE=\"${FLM_SOURCE_DIR}/include/AutoModel/all_models.hpp\"" + "FLM_TEST_MODEL_LIST_PATH=\"${FLM_SOURCE_DIR}/model_list.json\"") + target_compile_options(test_gemma4_registry PRIVATE + $<$:/wd4005 /wd4244>) + target_link_directories(test_gemma4_registry PRIVATE "${XRT_INCLUDE_DIR}/../lib") + target_link_libraries(test_gemma4_registry PRIVATE xrt_coreutil) + add_test(NAME test_gemma4_registry COMMAND test_gemma4_registry) +else() + message(WARNING "gemma4_rai: test_gemma4_registry will NOT be built -- " + "RYZENAI_CORELIB_INCLUDE_DIR=${RYZENAI_CORELIB_INCLUDE_DIR}, " + "BOOST_INCLUDE_DIR=${BOOST_INCLUDE_DIR}, XRT_INCLUDE_DIR=${XRT_INCLUDE_DIR}") + add_test(NAME test_gemma4_registry_WAS_NOT_BUILT + COMMAND ${CMAKE_COMMAND} -E echo "test_gemma4_registry WAS NOT BUILT") + set_tests_properties(test_gemma4_registry_WAS_NOT_BUILT PROPERTIES + FAIL_REGULAR_EXPRESSION "WAS NOT BUILT") endif() diff --git a/src/test/gemma4_rai/fake_corelib.cpp b/src/test/gemma4_rai/fake_corelib.cpp deleted file mode 100644 index e9a3e5613..000000000 --- a/src/test/gemma4_rai/fake_corelib.cpp +++ /dev/null @@ -1,1252 +0,0 @@ -#include "fake_corelib.hpp" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// See fake_corelib.hpp for why this is a separate, independent fake rather -// than a reuse of src/test/phi4_rai/fake_corelib.{hpp,cpp}, for the list of -// what it validates, and -- more importantly -- for the list of what it -// still cannot catch. This file implements every entry -// `FLM_CORELIB_FUNCTIONS` declares (corelib_api.hpp's constructor refuses to -// resolve if even one is missing). - -namespace { - -fake_corelib::State state; - -/// \brief serializes every entry point in this file -/// \note NOT decoration. C7's engine packs its weights across a pool of -/// `kWeightCreateConcurrency` threads, so eight of these functions run -/// at once and every one of them pushes onto a `State` vector. Held for -/// the whole body of each entry point, the same way -/// src/test/phi4_rai/fake_corelib.cpp holds its own. Recursive because -/// nothing here promises one entry point never calls another. -std::recursive_mutex state_mutex; - -/// \brief what `get_last_error_message` reports; set only by refusals that -/// carry corelib's own wording -std::string last_error; - -/// \brief the backing allocation a tensor and all its windows share -/// -/// Held by shared_ptr so that a window (and a window of a window) keeps the -/// root alive and, crucially, compares EQUAL to its siblings -- which is what -/// makes an aliasing check possible at all. See StorageOf. -struct FakeStorage { - std::size_t byte_size{}; -}; - -struct FakeObject { - std::string kind; - ryzenai_corelib_data_type data_type{ryzenai_corelib_data_type_bf16}; - std::vector shape; - std::size_t byte_size{}; - /// \brief offset from the ROOT allocation, in elements, accumulated - /// through however many levels of window - std::size_t window_offset{}; - /// \brief the strides a window was created with, empty for dense - std::vector strides; - std::shared_ptr storage; - /// \brief the packed N a matmul weights object was created for - /// \note What lets a matmul dispatch pick the SAME `"matmul-"` helper - /// key `matmul_pad_shape_tag` used to compute its bucket answers. - /// Without it a dispatch has no way to know which grid its own - /// operands should be checked against. - std::int64_t weight_n{-1}; - /// \brief how many bytes `weights_copy_data` reports for this object - /// \note Zero for anything that is not a packed weight. See - /// PackedSizeFor for why the number is small and synthetic. - std::size_t packed_size{0}; - /// \brief the ssmlp descriptor an ssmlp weights object was created with - /// \note `ryzenai_corelib_ssmlp_bf16` is handed a packed weights object - /// and no descriptor, so this is the only place the descriptor's - /// fields still exist by the time the dispatch happens. Carried - /// onto DispatchRecord so C7/C8 can assert on them. - ryzenai_corelib_ssmlp_bf16_weights_desc ssmlp_desc{}; - bool has_ssmlp_desc{false}; - /// \brief four bytes naming WHAT was packed into this weights object - /// - /// THE PACKED BYTES HAVE TO CARRY IT OR THE CACHE PATH CANNOT BE CHECKED - /// AT ALL. `weights_copy_data` stamps this at offset 0 of the blob it - /// hands out, so the weight cache's file contains it, and each - /// `..._weights_create_from_file` reads it back out of the slice it was - /// given. Without that, a cache-bound engine carries no identity anywhere - /// -- the descriptor says "some weight of this shape" and corelib returns - /// no components -- and `BindFromCache`'s assignment BY SLOT INDEX has - /// nothing watching it. - /// - /// This also makes the fake MORE faithful, not less: real corelib's - /// packed bytes are a function of the weight that was packed, so two - /// distinct tensors never produce one blob. Before this the bytes were a - /// function of the descriptor alone, which is a fidelity real corelib - /// does not have and a blindness a test inherits. - std::uint32_t identity_tag{0}; -}; - -void* NewObject(std::string kind = "generic") { - auto* object = new FakeObject; - object->kind = std::move(kind); - return object; -} - -/// \brief the synthetic packed size for one weight kind and shape -/// \note SMALL AND SYNTHETIC ON PURPOSE -- see fake_corelib.hpp's -/// PlePackedSize. The weight cache writes every packed weight to a -/// file, and an E2B model has 272 of them; reporting realistic sizes -/// would make a bookkeeping test write gigabytes. The only properties -/// that matter are determinism and that two different descriptors -/// disagree, so a cache slice written for one weight is refused for -/// another -- which is the thing corelib really does -/// ("corelib rejects a slice that is not exactly what the descriptor -/// packs to"). -std::size_t PackedSizeFor(std::uint32_t kind, std::int64_t k, std::int64_t n, - std::uint32_t group) { - return static_cast(64 + 11 * kind + (k % 251) + 3 * (n % 241) + - 7 * group); -} - -std::size_t Elements(const std::vector& shape) { - std::size_t result = 1; - for (const auto dimension : shape) result *= static_cast(dimension); - return result; -} - -std::size_t TypeBytes(ryzenai_corelib_data_type type) { - return RYZENAI_CORELIB_DATA_TYPE_BITS(type) / 8; -} - -/// \brief the fixture's four-byte tag at the head of a block stream, or 0 -/// \note WHAT MAKES ONE QUANTIZED TENSOR TELLABLE FROM ANOTHER. corelib's -/// requantizing entry points take a bare `const void*`; the descriptor -/// carries only shape and group size, and Gemma 4 ships many tensors -/// that agree on both (every layer's `attn_q`, and `ffn_gate` against -/// `ffn_up` on the same layer). See gemma4_gguf_fixture.hpp's -/// BlockSignatureAt for what is actually in those four bytes and why it -/// is not a valid Q8_0 block. -std::uint32_t BlockTag(const void* blocks) { - if (!blocks) return 0; - std::uint32_t tag = 0; - std::memcpy(&tag, blocks, sizeof(tag)); - return tag; -} - -/// \brief the identity tag this fake stamped into a packed blob, read back -/// from the slice a `..._weights_create_from_file` was pointed at -/// \return 0 when the file cannot be opened or the slice is too short -- -/// which is a MISS, not an error: every caller records it and the -/// assertion belongs to the test, exactly as on the pack leg where a -/// never-stamped payload also reads 0. -/// \note THE FAKE REALLY OPENS THE FILE. It has to: the whole point is that -/// the identity travelled through the weight cache's bytes on disk -/// rather than through a label this process kept in memory, which is -/// the same reason DispatchRecord::weights joins on the object that -/// crossed the ABI instead of on a name. Real corelib maps that slice; -/// reading four bytes of it is the cheapest honest imitation. -std::uint32_t IdentityTagInFile(const char* path, std::uint64_t offset, - std::uint64_t size) { - if (!path || size < sizeof(std::uint32_t)) return 0; - std::ifstream input(path, std::ios::binary); - if (!input) return 0; - input.seekg(static_cast(offset)); - std::uint32_t tag = 0; - input.read(reinterpret_cast(&tag), sizeof(tag)); - return input ? tag : 0; -} - -/// \brief round `rows` up the way every corelib pad helper here rounds: -/// unrounded at the M=1 decode kernel, to the next multiple of -/// `State::pad_multiple` otherwise, recording the answer under -/// `helper` as it goes -/// \note This is a stand-in for corelib's OWN grid, which differs per op and -/// per shape (see corelib.h's own per-op comments). Nothing in -/// test_gemma4_shape_plan.cpp asserts an exact padded value for -/// matmul/ssmlp/flat_mha rows, so one simple consistent rule suffices -/// -- but see fake_corelib.hpp: a bug that only shows up when two ops' -/// real grids DIVERGE cannot be seen through a single shared rule. -std::int64_t PaddedRows(std::string_view helper, std::int64_t rows) { - const auto multiple = state.pad_multiple; - const std::int64_t result = - (rows <= 1 || multiple <= 0) ? rows : (rows + multiple - 1) / multiple * multiple; - state.pad_answers_seen[std::string(helper)].insert(result); - return result; -} - -/// \brief the `"matmul-"` helper key an output width maps to -/// \note One key PER DISTINCT N, so a Q-projection operand is never checked -/// against the output projection's grid and vice versa. Keying every -/// matmul on one name would make OnGrid accept any matmul answer for -/// any matmul operand, which is most of what it is for. -std::string MatmulHelperFor(std::int64_t n) { - return "matmul-" + std::to_string(n); -} - -std::int64_t RowsOf(void* object) { - if (!object) return -1; - const auto& shape = static_cast(object)->shape; - return shape.empty() ? -1 : shape[0]; -} - -std::size_t RankOf(void* object) { - return object ? static_cast(object)->shape.size() : 0; -} - -/// \brief `object`'s declared shape, empty for a null or unshaped object -/// \note What `DispatchRecord::operand_shapes` carries. See that field for -/// why the extent beyond the leading one has to be retained at all. -std::vector ShapeOf(void* object) { - if (!object) return {}; - return static_cast(object)->shape; -} - -/// \brief the underlying allocation `object` reads and writes through -/// \note Two DIFFERENT FakeObjects can share this: every -/// `create_tensor_window` call mints a fresh object, but windows over -/// the same parent (however many levels deep) share that parent's -/// `storage` shared_ptr. This is what -/// DispatchRecord::operand_storage compares -- pointer identity of -/// `operands` itself only tells two OBJECTS apart, never two VIEWS of -/// the same memory apart. -const void* StorageOf(void* object) { - return object ? static_cast(object)->storage.get() : nullptr; -} - -/// \brief every rank-2 operand among these reports the SAME row count -/// \note THIS IS NOT THE CHECK REAL CORELIB MAKES -- corelib checks each -/// operand against ITS OWN pad answer; it has no notion of "these -/// operands must agree with each other" as a rule in itself. Comparing -/// operands pairwise is necessary (a dispatch with disagreeing operands -/// cannot be consistent with any single pad answer) but NOT sufficient: -/// operands that agree with each other while being uniformly wrong sail -/// through it, and OnGrid is what is meant to narrow that -- see -/// fake_corelib.hpp for how far it actually gets. A rank != 2 operand -/// (a KV cache's [heads, seq, head_dim]) is skipped rather than -/// compared: its row-count concept is unrelated to the others'. -bool RowsAgree(std::initializer_list operands) { - std::int64_t common = -1; - for (void* object : operands) { - if (!object || RankOf(object) != 2) continue; - const auto rows = RowsOf(object); - if (common == -1) common = rows; - else if (common != rows) return false; - } - return true; -} - -/// \brief `object`'s row count is a value `helper` is known to have produced -/// for SOME bucket -/// \return true (skipped, not failed) for a rank != 2 or null object -- -/// nothing to check -- or for an EMPTY helper name, meaning the -/// caller could not identify a helper at all (e.g. a matmul -/// dispatched with no real weights object) and has nothing to check -/// against. -/// \return **false** for a NAMED helper this engine never queried a pad -/// answer for -- i.e. an unqueried helper is a REJECTION here, not a -/// skip. -/// -/// THAT RULE IS ONLY AS GOOD AS THE PLAN'S COVERAGE, AND THE COVERAGE IS NOT -/// AUTOMATIC. An earlier version of this comment claimed "every helper name -/// in play is queried, and therefore recorded, before any dispatch can run, -/// because Gemma4ShapePlan::Build() enumerates every bucket at construction". -/// That is false and was demonstrated false by execution: `Build()` -/// enumerates every row BUCKET, but the matmul family is keyed per output -/// width (`MatmulHelperFor`), and a width `Build()` never queries has no -/// entry at all. Before this round `Build()` did not query lm_head, so C7's -/// very first dispatch -- 1 x hidden -> vocab -- came back `bad_argument` -/// from right here. A width the plan misses shows up as a FALSE REJECTION of -/// a legitimate dispatch, not as a missing check. -/// -/// The rejection is kept anyway, and deliberately: a silent skip would make -/// `OnGrid` useless for exactly the ops it matters most for, and a false -/// rejection is loud, traceable to this function, and fixed by planning the -/// missing width -- which is where it belongs. `Build()` now queries -/// lm_head for that reason. `"ssmlp"` and `"mha"` are single constant keys -/// `Build()` always registers, so only the matmul family can be caught out -/// this way. -/// -/// Conflating "unqueried" with the empty-name case above would hide both. An -/// equivalent fake-hardening test passed for the wrong reason on the Phi-4 -/// side once already -- keep the two branches, and this distinction, intact. -bool OnGrid(std::string_view helper, void* object) { - if (helper.empty() || !object || RankOf(object) != 2) return true; - const auto seen = state.pad_answers_seen.find(std::string(helper)); - if (seen == state.pad_answers_seen.end()) return false; - return seen->second.contains(RowsOf(object)); -} - -/// \brief `object`'s row count is one `ryzenai_corelib_ple_bf16` ships a -/// kernel for -/// \return true (skipped) for a null or rank != 2 object -/// -/// THE ONE OP WITH A KNOWABLE GRID AND NO PAD HELPER TO ASK. There is no -/// `ryzenai_corelib_ple_bf16_pad_rows` in the pinned 0.5.0 header, so -/// `pad_answers_seen` can never hold a ple entry and `OnGrid` would reject -/// every ple dispatch outright. What the header does give is the table -/// itself, verbatim, in `ryzenai_corelib_ple_bf16`'s own comment: M "rounds -/// up to a row count the ELF set ships -- 1 for decode, then 64, 128, 256, -/// 512, 1024, 2048, 3072, 4096 -- and the kernel then reads and writes that -/// many rows in full. So EVERY tensor must hold the padded extent, not the -/// live one; a buffer sized to the live rows is run off the end and comes -/// back with plausible numbers." -/// -/// STRICTER THAN CORELIB, DELIBERATELY. Corelib would accept a non-bucket M, -/// round it up and overrun. This refuses, so the overrun is a test failure -/// rather than a plausible one -- the same trade `OnGrid` makes for matmul. -/// -/// A SECOND, INDEPENDENT COPY of the table in -/// gemma4_rai_shape_plan.cpp's `kExecutionRows`, and it stays a second copy: -/// a fake that read the plan's table could not disagree with the plan. -bool OnPleGrid(void* object) { - if (!object || RankOf(object) != 2) return true; - static constexpr std::int64_t kPleRows[] = {1, 64, 128, 256, 512, - 1024, 2048, 3072, 4096}; - const auto rows = RowsOf(object); - for (const std::int64_t bucket : kPleRows) { - if (rows == bucket) return true; - } - return false; -} - -/// \brief true iff every one of `operands` that has storage at all has -/// DISTINCT storage from every other -- i.e. no two are windows -/// (however many levels deep) over the SAME backing allocation -/// \note corelib.h states this for `ssmlp_bf16`'s four tensors ("All four are -/// separate buffers; nothing is packed or interleaved"), and it is MORE -/// load-bearing on Gemma 4 than on Phi-4, not less: with -/// `post_feedforward_layernorm` set the kernel never writes the -/// `normalized` plane at all (corelib.h:1117-1128), yet still requires -/// a distinct buffer for that port. A port that is write-only in name -/// is exactly the one a reader is tempted to alias onto another. -/// Enforced HERE, at dispatch, not merely recorded for a caller to -/// notice afterward. Compares StorageOf, never the handle. A null or -/// storage-less operand is skipped rather than treated as "distinct -/// from everything", since there is nothing real to compare. -bool AllDistinctStorage(std::initializer_list operands) { - std::vector seen; - for (void* object : operands) { - const void* storage = StorageOf(object); - if (!storage) continue; - for (const void* other : seen) { - if (other == storage) return false; - } - seen.push_back(storage); - } - return true; -} - -/// \brief a (count, offset) element range lies inside `object`'s OWN extent -/// \note corelib.h: "The tensor's own size, which is what write/read are -/// bounded by. The underlying allocation is page-rounded and may be -/// larger." So this is checked against the declared shape, not against -/// the root allocation -- a window that reads past its own end but -/// stays inside the parent buffer is still a bug, and bounding against -/// the allocation would miss it. -bool WithinExtent(void* object, std::size_t count, std::size_t offset) { - if (!object) return false; - const auto* tensor = static_cast(object); - if (tensor->shape.empty()) return true; // an unshaped object bounds nothing - return offset + count <= Elements(tensor->shape); -} - -#define FLM_DEFINE_FAKE_TAG(member, symbol) \ - struct member##_tag { \ - static constexpr std::string_view name = #symbol; \ - }; -FLM_CORELIB_FUNCTIONS(FLM_DEFINE_FAKE_TAG) -#undef FLM_DEFINE_FAKE_TAG - -// The pad helpers are gone from corelib. The branches below still name their -// tags; they are never selected, and defining the types keeps those branches -// from being ill-formed. -struct matmul_pad_shape_tag {}; -struct ssmlp_pad_rows_tag {}; -struct flat_mha_pad_rows_tag {}; - -/// \brief the source with `role` in a create's flat source array, or NULL -const ryzenai_corelib_weights_source* FindSource(const ryzenai_corelib_weights_source* sources, - std::size_t count, - ryzenai_corelib_weights_role role) { - for (std::size_t i = 0; sources && i < count; ++i) - if (sources[i].role == role) return &sources[i]; - return nullptr; -} - -const void* SourceData(const ryzenai_corelib_weights_source* sources, std::size_t count, - ryzenai_corelib_weights_role role) { - const auto* source = FindSource(sources, count, role); - return source ? source->data : nullptr; -} - -template -inline constexpr bool kAlwaysFalse = false; - -template -struct TypedFake; - -template -struct TypedFake { - static Result Invoke(Args... args) { - // Held for the whole body. The engine's weight-create pool calls - // eight of these at once and every branch below touches `state`. - std::lock_guard state_lock(state_mutex); - auto arguments = std::forward_as_tuple(args...); - - if constexpr (std::is_same_v) { - if (std::get<0>(arguments)) *std::get<0>(arguments) = RYZENAI_CORELIB_VERSION_MAJOR; - if (std::get<1>(arguments)) *std::get<1>(arguments) = RYZENAI_CORELIB_VERSION_MINOR; - if (std::get<2>(arguments)) *std::get<2>(arguments) = RYZENAI_CORELIB_VERSION_PATCH; - return; - } else if constexpr (std::is_same_v) { - return "success"; - } else if constexpr (std::is_same_v) { - return last_error.c_str(); - } else if constexpr (std::is_same_v) { - static const char kKernelsRoot[] = ""; - return kKernelsRoot; - } else if constexpr (std::is_same_v) { - static const int kFakeDevice = 0; - return static_cast(&kFakeDevice); - } else if constexpr (std::is_same_v) { - // RECORD THE KIND BEFORE THE DELETE -- the string lives in the - // object being freed. See State::release_order for why the order - // of these calls is the only observable this fault has. - auto* object = static_cast(std::get<0>(arguments)); - if (object) state.release_order.push_back(object->kind); - delete object; - return; - } else if constexpr (std::is_same_v) { - // RECORDED INTO THE SAME LOG AS THE RELEASES. `cleanup` is the - // runtime's own teardown call -- `CorelibRuntime::ShutdownProcess` - // makes it and then drops the runtime -- so "the runtime goes - // last of all" becomes one question about one ABI sequence - // rather than two questions about two instruments. - state.release_order.push_back("cleanup"); - return; - } else if constexpr (std::is_same_v) { - // 0.5.0: (prefill_pdi, token_pdi, out). The shape plan itself - // never calls this -- MakeStreamForTest() hands back a stream - // directly -- but corelib_api.hpp still requires the symbol to - // resolve. - auto* out = std::get<2>(arguments); - if (out) *out = NewObject("stream"); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - const auto type = std::get<0>(arguments); - const auto* shape = std::get<1>(arguments); - const auto shape_len = std::get<2>(arguments); - auto* out = std::get<3>(arguments); - if (out) *out = nullptr; - if (!out || !shape) return ryzenai_corelib_status_bad_argument; - auto* object = static_cast(NewObject("tensor")); - object->data_type = type; - object->shape.assign(shape, shape + shape_len); - object->byte_size = Elements(object->shape) * TypeBytes(type); - object->storage = std::make_shared(); - object->storage->byte_size = object->byte_size; - *out = object; - state.tensor_creates.push_back({type, object->shape, object}); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - const auto type = std::get<0>(arguments); - const auto* shape = std::get<1>(arguments); - const auto shape_len = std::get<2>(arguments); - const void* data = std::get<3>(arguments); - auto* out = std::get<4>(arguments); - if (out) *out = nullptr; - if (!out || !shape || !data) return ryzenai_corelib_status_bad_argument; - auto* object = static_cast(NewObject("host_view")); - object->data_type = type; - object->shape.assign(shape, shape + shape_len); - object->byte_size = Elements(object->shape) * TypeBytes(type); - *out = object; - state.host_view_creates.push_back({type, object->shape, object}); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - const char* model = std::get<1>(arguments); - state.kernel_bundles.push_back( - {std::get<0>(arguments), model ? model : "", state.enum_calls.size()}); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - void* parent = std::get<0>(arguments); - const auto* shape = std::get<1>(arguments); - const auto shape_len = std::get<2>(arguments); - const auto offset = std::get<3>(arguments); - auto* out = std::get<6>(arguments); - if (out) *out = nullptr; - if (!out || !shape) return ryzenai_corelib_status_bad_argument; - ryzenai_corelib_data_type data_type = ryzenai_corelib_data_type_bf16; - std::shared_ptr storage; - std::size_t window_offset = offset; - if (parent) { - const auto* parent_object = static_cast(parent); - data_type = parent_object->data_type; - storage = parent_object->storage; - window_offset = parent_object->window_offset + offset; - } - const std::vector window_shape(shape, shape + shape_len); - // A WINDOW MUST FIT INSIDE THE ALLOCATION IT IS CARVED FROM. - // Checked in ELEMENTS of the (possibly nested) parent's own - // dtype, against the ROOT allocation's real byte size, which - // every window in a chain shares via `storage` regardless of how - // many levels deep it is. Gemma 4 needs this specifically: its - // activation buffers are allocated once at the WIDER geometry and - // windowed down per layer on the column axis, so an off-by-one in - // that arithmetic is a live failure mode and not a hypothetical. - const auto requested_elements = window_offset + Elements(window_shape); - if (storage && requested_elements * TypeBytes(data_type) > storage->byte_size) { - return ryzenai_corelib_status_bad_argument; - } - auto* object = static_cast(NewObject("window")); - object->data_type = data_type; - object->storage = storage; - object->window_offset = window_offset; - object->shape = window_shape; - object->byte_size = Elements(object->shape) * TypeBytes(object->data_type); - const auto* strides = std::get<4>(arguments); - const auto strides_len = std::get<5>(arguments); - if (strides) object->strides.assign(strides, strides + strides_len); - *out = object; - state.tensor_windows.push_back( - {parent, object->shape, offset, object, object->strides}); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - auto* object = static_cast(std::get<0>(arguments)); - if (std::get<1>(arguments)) - *std::get<1>(arguments) = object ? object->byte_size : 0; - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - auto* object = static_cast(std::get<0>(arguments)); - if (std::get<1>(arguments)) - *std::get<1>(arguments) = - object ? object->data_type : ryzenai_corelib_data_type_bf16; - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - void* tensor = std::get<0>(arguments); - const auto count = std::get<3>(arguments); - const auto offset = std::get<4>(arguments); - const bool accepted = WithinExtent(tensor, count, offset); - state.tensor_writes.push_back( - {tensor, std::get<1>(arguments), count, offset, accepted}); - // No bytes are moved: nothing in src/test/gemma4_rai reads a - // tensor's contents back, and a fake that pretended to would be - // claiming a fidelity it does not have. The BOUND is the part - // that matters, and it is enforced. - return accepted ? ryzenai_corelib_status_success - : ryzenai_corelib_status_bad_argument; - } else if constexpr (std::is_same_v) { - void* tensor = std::get<0>(arguments); - const auto destination_type = std::get<1>(arguments); - void* destination = std::get<2>(arguments); - const auto count = std::get<3>(arguments); - const auto offset = std::get<4>(arguments); - const bool accepted = WithinExtent(tensor, count, offset); - // RECORDED, INCLUDING WHEN REFUSED. See TensorReadRecord: the - // (count, offset) pair is what "read the wrong row" is made of, - // and it is the only part of a read this fake can honestly answer - // for. NO DEVICE BYTES ARE MOVED -- there are none -- so an - // ACCEPTED read fills the destination with a synthetic constant - // instead (below), and a REFUSED one returns before writing - // anything at all. - state.tensor_reads.push_back( - {tensor, destination_type, count, offset, accepted}); - if (!accepted) return ryzenai_corelib_status_bad_argument; - // A SYNTHETIC CONSTANT, NOT ZERO. See fake_corelib.hpp's - // kSyntheticReadValue: zeros are a fixed point of the host-side - // logit softcap, so a read that returned them made "the engine - // applied the cap" and "the engine dropped it" the same - // observation. Written in the CALLER's dtype, because that is what - // `count` is in. - // - // ... UNLESS A TEST ASKED FOR A NaN. See - // State::read_returns_nonfinite: the engine's RequireFinite guard - // is otherwise unreachable, and an unreachable guard is one - // nothing says is still wired up. - // - // AND, OPTIONALLY, ONLY ON ONE READ. See - // State::nonfinite_read_index: the engine checks TWICE, and a - // global injection always trips the first check and never the - // second. The record for the read in progress has just been - // pushed, so its index is `tensor_reads.size() - 1`. - const bool nonfinite = - state.read_returns_nonfinite && - (!state.nonfinite_read_index || - *state.nonfinite_read_index == state.tensor_reads.size() - 1); - const float value = nonfinite - ? std::numeric_limits::quiet_NaN() - : fake_corelib::kSyntheticReadValue; - if (destination) { - if (destination_type == ryzenai_corelib_data_type_fp32) { - auto* out = static_cast(destination); - for (std::size_t i = 0; i < count; ++i) out[i] = value; - } else if (destination_type == ryzenai_corelib_data_type_bf16) { - const auto bits = std::bit_cast(value); - const auto as_bf16 = static_cast(bits >> 16); - auto* out = static_cast(destination); - for (std::size_t i = 0; i < count; ++i) out[i] = as_bf16; - } else { - std::memset(destination, 0, count * TypeBytes(destination_type)); - } - } - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - // A context above one page plans against this grid, not - // flat_mha's, so an empty answer leaves every prefill bucket - // without a kernel. - auto fn = std::get<1>(arguments); - if (fn && state.shipped_paged.empty()) { - fn(std::get<2>(arguments), 0, 0, 0, -1, 0, 0, 0, 0); - } else if (fn) { - for (const auto& shipped : state.shipped_paged) { - fn(std::get<2>(arguments), 0, shipped.heads, shipped.kv_heads, shipped.m, 0, - shipped.head_size, shipped.page_size, 0); - } - } - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - // 0.5.0: (stream, m*, k*, n*, group_size). Every padding helper - // gained a leading stream in 0.5.0 -- recorded here so - // "the stream was passed to every pad helper" has something to - // check. - auto stream = std::get<0>(arguments); - auto* m = std::get<1>(arguments); - auto* k = std::get<2>(arguments); - auto* n = std::get<3>(arguments); - const auto group = std::get<4>(arguments); - // Recorded BEFORE any perturbation, so the record always states - // the logical shape the caller asked about. - state.matmul_pad_calls.push_back( - {stream, m ? *m : -1, k ? *k : -1, n ? *n : -1, group}); - if (m) *m = PaddedRows(MatmulHelperFor(n ? *n : -1), *m); - if (k) *k += state.matmul_k_delta; - if (n) *n += state.matmul_n_delta; - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - // (desc, sources, count, options, out). A lone `packed` source is - // a cache bind; anything else packs model components. - const auto* desc = std::get<0>(arguments); - const auto* sources = std::get<1>(arguments); - const auto count = std::get<2>(arguments); - auto* out = std::get<4>(arguments); - if (out) *out = nullptr; - if (!desc || !out) return ryzenai_corelib_status_bad_argument; - const auto packed = PackedSizeFor(1, desc->k, desc->n, desc->group_size); - const auto* blob = FindSource(sources, count, ryzenai_corelib_weights_role_packed); - const bool kFromFile = blob != nullptr; - std::uint32_t blocks_tag = 0; - if (kFromFile) { - // corelib.h: the length must be exactly what this descriptor - // packs to. A cache entry written for another weight is - // refused here rather than bound and believed. That check is - // blind to a swap between two weights of the SAME descriptor, - // which is why the identity comes out of the bytes. - const auto size = static_cast(blob->size); - if (count != 1 || size != packed) return ryzenai_corelib_status_bad_argument; - blocks_tag = IdentityTagInFile(blob->path, blob->offset, size); - } else { - const void* qweight = - SourceData(sources, count, ryzenai_corelib_weights_role_qweight); - if (!qweight) return ryzenai_corelib_status_bad_argument; - blocks_tag = BlockTag(qweight); - } - *out = NewObject("matmul_weights"); - auto* object = static_cast(*out); - object->weight_n = desc->n; - object->packed_size = packed; - static constexpr std::int64_t kBuckets[] = { - 1, 64, 128, 256, 512, 1024, 2048, 3072, 4096}; - for (const std::int64_t bucket : kBuckets) - PaddedRows(MatmulHelperFor(desc->n), bucket); - object->identity_tag = blocks_tag; - state.matmul_weights_creates.push_back({*desc, kFromFile, blocks_tag, *out}); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - // (desc, sources, count, options, out). - const auto* desc = std::get<0>(arguments); - const auto* sources = std::get<1>(arguments); - const auto count = std::get<2>(arguments); - auto* out = std::get<4>(arguments); - if (out) *out = nullptr; - if (!desc || !out) return ryzenai_corelib_status_bad_argument; - const auto packed = PackedSizeFor(3, desc->k, 0, 0); - if (const auto* blob = - FindSource(sources, count, ryzenai_corelib_weights_role_packed)) { - const auto size = static_cast(blob->size); - if (count != 1 || size != packed) return ryzenai_corelib_status_bad_argument; - // WHICH GAMMA THESE BYTES WERE PACKED FROM. corelib hands no - // components back on this leg, so `scale` stays empty and this - // tag -- recovered from the slice itself -- is the only thing - // that can tell one layer's Q-norm from its K-norm, whose - // descriptors and packed lengths are identical. - const auto tag = IdentityTagInFile(blob->path, blob->offset, size); - *out = NewObject("rmsnorm_weights"); - static_cast(*out)->packed_size = packed; - static_cast(*out)->identity_tag = tag; - state.rmsnorm_weights_creates.push_back( - {-1, *desc, true, {}, static_cast(tag), *out}); - return ryzenai_corelib_status_success; - } - const auto* scale = static_cast( - SourceData(sources, count, ryzenai_corelib_weights_role_scale)); - if (!scale) return ryzenai_corelib_status_bad_argument; - fake_corelib::RmsNormWeightsCreateCall call{0, *desc, false, {}, 0, nullptr}; - call.scale.assign(scale, scale + static_cast(desc->k)); - call.scale_tag = call.scale.empty() ? 0 : call.scale.front(); - *out = NewObject("rmsnorm_weights"); - static_cast(*out)->packed_size = packed; - static_cast(*out)->identity_tag = call.scale_tag; - call.object = *out; - state.rmsnorm_weights_creates.push_back(std::move(call)); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - // (desc, sources, count, options, out). A lone `packed` source is - // a cache bind and is recorded as a create alone; model - // components are packed AND created in this one call, so they - // land in both records. - const auto* desc = std::get<0>(arguments); - const auto* sources = std::get<1>(arguments); - const auto count = std::get<2>(arguments); - auto* out = std::get<4>(arguments); - if (out) *out = nullptr; - if (!desc || !out) return ryzenai_corelib_status_bad_argument; - if (desc->k <= 0 || desc->n <= 0) - return ryzenai_corelib_status_bad_argument; - // "32 IS THE ONLY VALUE mladfple ships, and anything else is - // refused by name" -- refused here too, so an engine that packed - // at `group` instead of `ple_group` could not pass by accident. - if (desc->group_size != 32) return ryzenai_corelib_status_unsupported; - const auto size = fake_corelib::PlePackedSize(*desc); - if (const auto* blob = - FindSource(sources, count, ryzenai_corelib_weights_role_packed)) { - const auto blob_size = static_cast(blob->size); - // corelib.h: "The length must be exactly what this descriptor - // packs to: a truncated blob is still a plausible one." - if (count != 1 || blob_size != size) { - state.ple_create_calls.push_back({*desc, blob_size, false, nullptr}); - return ryzenai_corelib_status_bad_argument; - } - *out = NewObject("ple_weights"); - static_cast(*out)->packed_size = size; - state.ple_create_calls.push_back({*desc, blob_size, true, *out}); - return ryzenai_corelib_status_success; - } - const auto* gate = static_cast( - SourceData(sources, count, ryzenai_corelib_weights_role_gate)); - const auto* proj = static_cast( - SourceData(sources, count, ryzenai_corelib_weights_role_proj)); - const auto* post_norm = static_cast( - SourceData(sources, count, ryzenai_corelib_weights_role_post_norm)); - const auto* next_norm = static_cast( - SourceData(sources, count, ryzenai_corelib_weights_role_next_norm)); - if (!gate || !proj || !post_norm || !next_norm) - return ryzenai_corelib_status_bad_argument; - fake_corelib::PlePackCall call{*desc, false, true, size, - 0.0f, 0.0f, 0.0f, 0.0f, - 0.0f, 0.0f, {}, {}}; - const auto k = static_cast(desc->k); - call.post_norm.assign(post_norm, post_norm + k); - call.next_norm.assign(next_norm, next_norm + k); - call.gate_first = gate[0]; - call.proj_first = proj[0]; - // THE ORIENTATION PROBES. gate is required as [k, n] row-major, - // so its (row 1, col 0) is at flat index n; proj is required as - // [n, k], so its (row 1, col 0) is at flat index k. Together with - // index 1 -- (row 0, col 1) in both -- these are the two - // positions a transpose exchanges. Index 0 and index k*n-1, the - // only two this fake used to read, are FIXED POINTS of a - // transpose and can never see one. See PlePackCall. - const auto n = static_cast(desc->n); - call.gate_at_1 = gate[1]; - call.gate_at_n = gate[n]; - call.proj_at_1 = proj[1]; - call.proj_at_k = proj[k]; - // Touch the far end of both matrices so an engine that passed a - // wrongly-sized buffer runs off it here rather than only on - // hardware: gate is [k, n] and proj is [n, k], both k*n floats. - const auto elements = k * n; - volatile float sink = 0.0f; - sink += gate[elements - 1]; - sink += proj[elements - 1]; - (void)sink; - state.ple_pack_calls.push_back(std::move(call)); - *out = NewObject("ple_weights"); - static_cast(*out)->packed_size = size; - state.ple_create_calls.push_back({*desc, size, true, *out}); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v || - std::is_same_v || - std::is_same_v) { - auto* callback = std::get<1>(arguments); - auto* ctx = std::get<2>(arguments); - state.enum_calls.push_back({std::string(Tag::name), std::get<0>(arguments)}); - if (callback && state.explicit_grid) { - if constexpr (std::is_same_v) { - for (const auto& s : state.shipped_matmul) - callback(ctx, 0, s.m, s.k, s.n, s.group, false, - static_cast(0), static_cast(0)); - } else if constexpr (std::is_same_v) { - for (const auto& s : state.shipped_ssmlp) - callback(ctx, 0, s.family.c_str(), s.m, s.k, s.n, s.group); - } else { - for (const auto& s : state.shipped_mha) - callback(ctx, 0, s.heads, s.kv_heads, s.m, - static_cast(0), s.head_size, s.max_seq, - static_cast(0), s.window, s.kv_shared, - s.scale_one); - } - } else if (callback) { - if constexpr (std::is_same_v) { - callback(ctx, 0, static_cast(-1), - static_cast(0), static_cast(0), - static_cast(0), false, - static_cast(0), static_cast(0)); - } else if constexpr (std::is_same_v) { - static constexpr std::int64_t kBuckets[] = { - 1, 64, 128, 256, 512, 1024, 2048, 3072, 4096}; - for (const std::int64_t bucket : kBuckets) - PaddedRows("ssmlp", bucket); - callback(ctx, 0, "gemma_fusion", static_cast(-1), - static_cast(0), static_cast(0), - static_cast(0)); - } else { - static constexpr std::int64_t kBuckets[] = { - 1, 64, 128, 256, 512, 1024, 2048, 3072, 4096}; - for (const std::int64_t bucket : kBuckets) - PaddedRows("mha", bucket); - callback(ctx, 0, static_cast(-1), - static_cast(0), static_cast(-1), - static_cast(0), static_cast(0), - static_cast(0), static_cast(0), - static_cast(0), false, false); - } - } - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - // 0.5.0: (stream, m*, desc*). The WHOLE descriptor is recorded, - // not a k/n/group_size projection of it -- see fake_corelib.hpp. - auto stream = std::get<0>(arguments); - auto* m = std::get<1>(arguments); - const auto* desc = std::get<2>(arguments); - state.rows_pad_calls.push_back( - {stream, m ? *m : -1, - desc ? *desc : ryzenai_corelib_ssmlp_bf16_weights_desc{}}); - if (m) *m = PaddedRows("ssmlp", *m); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - // (desc, sources, count, options, out). - const auto* desc = std::get<0>(arguments); - const auto* sources = std::get<1>(arguments); - const auto count = std::get<2>(arguments); - auto* out = std::get<4>(arguments); - if (out) *out = nullptr; - if (!desc || !out) return ryzenai_corelib_status_bad_argument; - const auto packed = PackedSizeFor(2, desc->k, desc->n, desc->group_size); - const auto* blob = FindSource(sources, count, ryzenai_corelib_weights_role_packed); - const bool kFromFile = blob != nullptr; - fake_corelib::SsMlpWeightsCreateCall call{ - *desc, kFromFile, false, 0, 0, 0, 0, 0, 0, nullptr}; - if (kFromFile) { - const auto size = static_cast(blob->size); - if (count != 1 || size != packed) return ryzenai_corelib_status_bad_argument; - // ONLY `gate_tag`, and only because a packed blob has no - // ports. On the cache leg there are no three block streams to - // exchange -- corelib is handed one opaque slice -- so the - // question is which LAYER'S ssmlp this is, not which stream - // reached which port. `up_tag`/`down_tag` stay 0 here and the - // port assignment stays the pack leg's business. - call.gate_tag = IdentityTagInFile(blob->path, blob->offset, size); - } else { - // The two norms are the fields Gemma 4 must get right and - // Phi-4's fake never recorded. norm0 is this layer's - // `ffn_norm` and norm1 its `post_ffw_norm` -- BOTH this - // layer's, unlike the silu family whose norm1 belongs to the - // next one (gemma4_driver.py:1441-1449). One BF16 element of - // each is enough to say WHICH tensor was passed, because the - // fixture fills every F32 tensor with a constant signature - // derived from its name. - const auto data = [&](ryzenai_corelib_weights_role role) { - return SourceData(sources, count, role); - }; - if (!data(ryzenai_corelib_weights_role_gate_qweight)) - return ryzenai_corelib_status_bad_argument; - call.has_components = true; - const auto first = [](const void* p) -> std::uint16_t { - return p ? *static_cast(p) : 0; - }; - call.epsilon_bf16 = first(data(ryzenai_corelib_weights_role_epsilon)); - call.norm0_bf16 = first(data(ryzenai_corelib_weights_role_norm0)); - call.norm1_bf16 = first(data(ryzenai_corelib_weights_role_norm1)); - // AND WHICH BLOCK STREAM REACHED WHICH OF THE THREE PORTS. - // `gate` and `up` have the same shape on every Gemma 4 layer, - // so the descriptor cannot tell them apart and neither can a - // length check. See SsMlpWeightsCreateCall::gate_tag. - call.gate_tag = BlockTag(data(ryzenai_corelib_weights_role_gate_qweight)); - call.up_tag = BlockTag(data(ryzenai_corelib_weights_role_up_qweight)); - call.down_tag = BlockTag(data(ryzenai_corelib_weights_role_down_qweight)); - } - *out = NewObject("ssmlp_weights"); - // The descriptor is kept HERE because the dispatch never sees - // one -- see FakeObject::ssmlp_desc. - auto* object = static_cast(*out); - object->has_ssmlp_desc = true; - object->ssmlp_desc = *desc; - object->packed_size = packed; - object->identity_tag = call.gate_tag; - call.object = *out; - state.ssmlp_weights_creates.push_back(call); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - // Two-call, like the real one: `out == NULL` reports the size. - // It used to report 0 unconditionally, which made every caller - // treat the copy as unavailable -- the weight cache's write path - // in particular gave up silently and nothing noticed. - auto* weights = std::get<0>(arguments); - void* out = std::get<1>(arguments); - const auto out_size = std::get<2>(arguments); - auto* size = std::get<3>(arguments); - if (!weights || !size) return ryzenai_corelib_status_bad_argument; - const auto* object = static_cast(weights); - const auto packed = object->packed_size; - *size = packed; - if (out) { - if (out_size < packed) return ryzenai_corelib_status_bad_argument; - std::memset(out, static_cast(packed & 0xFF), packed); - // THE BLOB NAMES WHAT WAS PACKED INTO IT, at offset 0, which - // is what puts the identity into the weight cache's file and - // therefore onto the cache path. See FakeObject::identity_tag. - // Filler first, then the tag, so the tag is never overwritten. - if (packed >= sizeof(std::uint32_t)) - std::memcpy(out, &object->identity_tag, sizeof(std::uint32_t)); - } - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - // 0.5.0: (stream, m*, desc*). - auto stream = std::get<0>(arguments); - auto* m = std::get<1>(arguments); - auto* desc = std::get<2>(arguments); - state.mha_pad_calls.push_back( - {stream, m ? *m : -1, - desc ? *desc : ryzenai_corelib_flat_mha_bf16_desc{}}); - if (m) *m = PaddedRows("mha", *m); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v || - std::is_same_v || - std::is_same_v || - std::is_same_v || - std::is_same_v || - std::is_same_v) { - // THE INJECTED MID-CHAIN REFUSAL, BEFORE ANY VALIDATION AND - // BEFORE THE RECORD. See State::fail_dispatch_at: this is the - // only way a test can reach the engine's catch block, which is - // where `poisoned` is set. Nothing is recorded for a refused - // call, so the index a test sets stays fixed. - if (state.fail_dispatch_at && - state.dispatches.size() == *state.fail_dispatch_at) { - return ryzenai_corelib_status_failure; - } - // Which operands are row-bearing differs by op: matmul's are - // input/output; ssmlp's are input/residual/skip_sum/normalized - // (its weights are not a row-shaped tensor); rmsnorm's are - // input/output; ple's are x/ple/out/norm_out; flat_mha's are - // q/k/out (cos/sin/caches have their own, deliberately different, - // extent rules and are not checked here). - // - // RowsAgree applies to every op, for a reason that is PARTIAL -- - // see RowsAgree and OnGrid above, and the "what it still cannot - // catch" list in fake_corelib.hpp. The rest is per-op, because - // corelib.h's requirements are: matmul/ssmlp/mha get OnGrid, - // ple gets OnPleGrid instead (it has no pad helper to have - // produced an answer), and rmsnorm gets NEITHER. ssmlp and ple - // each get a distinctness rule their own section states; matmul, - // mha and rmsnorm have none. - std::vector operands; - // Declared before the branches so each can record the descriptor - // it was actually handed -- see DispatchRecord::mha_desc for why - // the plan's own copy is not a substitute for it. - fake_corelib::DispatchRecord record{}; - if constexpr (std::is_same_v) { - void* input = std::get<1>(arguments); - void* weights = std::get<2>(arguments); - void* output = std::get<3>(arguments); - operands = {input, output}; - // THE OBJECT, NOT JUST ITS SHAPE. See DispatchRecord::weights. - record.weights = weights; - const auto weight_n = - weights ? static_cast(weights)->weight_n : -1; - // No weights object, or one created without an N: no helper - // can be identified, so OnGrid is skipped rather than guessed - // at. RowsAgree still applies. - const std::string helper = - weight_n < 0 ? std::string() : MatmulHelperFor(weight_n); - if (!RowsAgree({input, output}) || !OnGrid(helper, input) || - !OnGrid(helper, output)) - return ryzenai_corelib_status_bad_argument; - // Like corelib: a 3-D output is the cache-writing form and - // must be a strided window [k * head, head, 1] with k >= its - // rows. A dense window over the whole cache is refused too. - if (output && RankOf(output) == 3) { - const auto* object = static_cast(output); - const auto& strides = object->strides; - const auto head = object->shape[2]; - if (head <= 0 || strides.size() != 3 || strides[2] != 1 || strides[1] != head || - strides[0] % head != 0 || strides[0] / head < object->shape[1]) { - last_error = "a cache-writing 3-D output must be a window with " - "strides [max_seq * head_size, head_size, 1]"; - return ryzenai_corelib_status_bad_argument; - } - } - } else if constexpr (std::is_same_v) { - void* input = std::get<1>(arguments); - void* residual = std::get<2>(arguments); - void* weights = std::get<3>(arguments); - void* skip_sum = std::get<4>(arguments); - void* normalized = std::get<5>(arguments); - operands = {input, residual, skip_sum, normalized}; - // THE OBJECT, NOT JUST ITS SHAPE. See DispatchRecord::weights. - record.weights = weights; - // The descriptor travelled with the WEIGHTS, not with this - // call -- corelib.h:1384-1388 passes no descriptor here. - if (weights) { - const auto* object = static_cast(weights); - record.has_ssmlp_desc = object->has_ssmlp_desc; - record.ssmlp_desc = object->ssmlp_desc; - } - if (!RowsAgree({input, residual, skip_sum, normalized}) || - !OnGrid("ssmlp", input) || !OnGrid("ssmlp", residual) || - !OnGrid("ssmlp", skip_sum) || !OnGrid("ssmlp", normalized) || - !AllDistinctStorage({input, residual, skip_sum, normalized})) - return ryzenai_corelib_status_bad_argument; - } else if constexpr (std::is_same_v) { - // (stream, input, weights, output). - void* input = std::get<1>(arguments); - void* weights = std::get<2>(arguments); - void* output = std::get<3>(arguments); - operands = {input, output}; - record.weights = weights; - // NO DISTINCTNESS RULE, AND THAT IS NOT AN OVERSIGHT. - // corelib.h: input and output "may be the SAME tensor FOR - // SMALL M -- the kernel reads a row before it writes it, so - // normalizing in place is safe there and saves a buffer". A - // rule here would be a false rejection of a legal decode - // dispatch. The header also says in-place CORRUPTS at - // M = 1024 / 2048 at k=128 even where an EXACT kernel ships, - // which is a property of the width and the row count that - // nothing host-side can distinguish from the safe case -- - // so it is not checkable here at all, by this fake or any - // other. Hardware, or a separate output tensor. - // - // NO GRID CHECK EITHER. `rmsnorm_bf16_pad_rows` is "an - // IDENTITY FUNCTION today and does not report the covering - // extent", and it is not even bound in - // FLM_CORELIB_FUNCTIONS, so there is no answer set to check - // a row count against and no honest one to invent. What the - // header warns of -- the smallest covering kernel writing - // PAST the live rows -- is an allocation-extent question, and - // the buffer a dispatch is handed here is already whatever - // the caller allocated. Real hardware is the check. - if (!RowsAgree({input, output})) - return ryzenai_corelib_status_bad_argument; - } else if constexpr (std::is_same_v) { - // (stream, x, ple, weights, out, norm_out). - void* x = std::get<1>(arguments); - void* slice = std::get<2>(arguments); - void* weights = std::get<3>(arguments); - void* out = std::get<4>(arguments); - void* norm_out = std::get<5>(arguments); - operands = {x, slice, out, norm_out}; - record.weights = weights; - // corelib.h: "`out` and `norm_out` must be distinct buffers, - // from each other and from `x`." `ple` -- the per-layer input - // slice -- is NOT in that set and is not added to it; the - // header does not put it there and inventing a rule corelib - // does not have is a false rejection waiting for C8. - if (!RowsAgree({x, slice, out, norm_out}) || !OnPleGrid(x) || - !AllDistinctStorage({x, out, norm_out})) - return ryzenai_corelib_status_bad_argument; - } else { - const auto* desc = std::get<1>(arguments); - void* q = std::get<2>(arguments); - void* k = std::get<3>(arguments); - void* out = std::get<9>(arguments); - operands = {q, k, out}; - // THE FIELD C7/C8 MUST BE ABLE TO ASSERT ON. Nothing here - // judges it -- neither this fake nor corelib can tell an - // owning dispatch handed a kvshare descriptor from a correct - // one -- but recording it is what makes the reference - // driver's load-time guard (gemma4_driver.py:1670-1690) - // portable to a test. - if constexpr (std::is_same_v) { - record.has_mha_desc = desc != nullptr; - if (desc) record.mha_desc = *desc; - } - // THE ROTARY PAIR AND THE TWO CACHES. Not row-bearing, so not - // in `operands` and not checked here -- retained so that "this - // layer roped with the sliding table and attended over cache - // 13" is a question a caller can ask at all. See - // DispatchRecord::mha_cos. - record.mha_position = std::get<4>(arguments); - record.mha_cos = std::get<5>(arguments); - record.mha_sin = std::get<6>(arguments); - record.mha_k_cache = std::get<7>(arguments); - record.mha_v_cache = std::get<8>(arguments); - // Like corelib: a rotary table that is not a host view is refused. - for (void* table : {record.mha_cos, record.mha_sin}) { - if (!table || static_cast(table)->kind != "host_view") - return ryzenai_corelib_status_bad_argument; - } - // paged_mha has no pad helper, so there is no answer set to - // put its rows on. - if (!RowsAgree({q, k, out})) - return ryzenai_corelib_status_bad_argument; - if constexpr (std::is_same_v) { - if (!OnGrid("mha", q) || !OnGrid("mha", k) || !OnGrid("mha", out)) - return ryzenai_corelib_status_bad_argument; - } - } - record.kind = std::is_same_v ? "matmul" - : std::is_same_v ? "ssmlp" - : std::is_same_v ? "rmsnorm" - : std::is_same_v ? "ple" - : std::is_same_v ? "paged_mha" - : "mha"; - record.stream = std::get<0>(arguments); - record.operands = operands; - record.operand_storage.reserve(operands.size()); - record.operand_shapes.reserve(operands.size()); - for (void* operand : operands) { - record.operand_storage.push_back(StorageOf(operand)); - // THE WIDTH, NOT ONLY THE ROWS. See - // DispatchRecord::operand_shapes: everything else this record - // holds is invariant under binding the other geometry's - // windows. - record.operand_shapes.push_back(ShapeOf(operand)); - } - // 0.5.0 removed the row count from every dispatch: M is the - // leading extent of the operand that was bound. - record.rows = operands.empty() ? -1 : RowsOf(operands.front()); - state.dispatches.push_back(std::move(record)); - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - return ryzenai_corelib_status_success; - } else if constexpr (std::is_same_v) { - return ryzenai_corelib_status_success; - } else { - static_assert(kAlwaysFalse, "unhandled fake corelib ABI result"); - } - } -}; - -#define FLM_ASSERT_FAKE_ABI(member, symbol) \ - static_assert(std::is_same_v< \ - decltype(&TypedFake::Invoke), \ - decltype(&::symbol)>); -FLM_CORELIB_FUNCTIONS(FLM_ASSERT_FAKE_ABI) -#undef FLM_ASSERT_FAKE_ABI - -void* FunctionFor(std::string_view name) { -#define FLM_MAP_FAKE_FUNCTION(member, symbol) \ - if (name == #symbol) { \ - return reinterpret_cast( \ - &TypedFake::Invoke); \ - } - FLM_CORELIB_FUNCTIONS(FLM_MAP_FAKE_FUNCTION) -#undef FLM_MAP_FAKE_FUNCTION - return nullptr; -} - -} // namespace - -namespace fake_corelib { - -State& GetState() { return state; } - -std::size_t PlePackedSize(const ryzenai_corelib_ple_bf16_weights_desc& desc) { - return PackedSizeFor(4, desc.k, desc.n, desc.group_size); -} - -ryzenai_corelib_status NoteMatmulPad(ryzenai_corelib_stream_ptr stream, - std::int64_t* rows, - std::int64_t* k, - std::int64_t* n, - std::int64_t group) { - std::lock_guard lock(state_mutex); - state.matmul_pad_calls.push_back( - {stream, rows ? *rows : -1, k ? *k : -1, n ? *n : -1, - static_cast(group)}); - if (rows) *rows = PaddedRows(MatmulHelperFor(n ? *n : -1), *rows); - if (k) *k += state.matmul_k_delta; - if (n) *n += state.matmul_n_delta; - return ryzenai_corelib_status_success; -} - -ryzenai_corelib_status NoteSsmlpRows( - ryzenai_corelib_stream_ptr stream, std::int64_t* rows, - const ryzenai_corelib_ssmlp_bf16_weights_desc* desc) { - std::lock_guard lock(state_mutex); - state.rows_pad_calls.push_back( - {stream, rows ? *rows : -1, - desc ? *desc : ryzenai_corelib_ssmlp_bf16_weights_desc{}}); - if (rows) *rows = PaddedRows("ssmlp", *rows); - return ryzenai_corelib_status_success; -} - -ryzenai_corelib_status NoteMhaRows(ryzenai_corelib_stream_ptr stream, - std::int64_t* rows, - const ryzenai_corelib_flat_mha_bf16_desc* desc) { - std::lock_guard lock(state_mutex); - state.mha_pad_calls.push_back( - {stream, rows ? *rows : -1, - desc ? *desc : ryzenai_corelib_flat_mha_bf16_desc{}}); - if (rows) *rows = PaddedRows("mha", *rows); - return ryzenai_corelib_status_success; -} - -void Reset() { - std::lock_guard lock(state_mutex); - state.enum_calls.clear(); - state.kernel_bundles.clear(); - state.explicit_grid = false; - state.shipped_matmul.clear(); - state.shipped_ssmlp.clear(); - state.shipped_mha.clear(); - state.shipped_paged.clear(); - state.matmul_pad_calls.clear(); - state.rows_pad_calls.clear(); - state.mha_pad_calls.clear(); - state.matmul_weights_creates.clear(); - state.ssmlp_weights_creates.clear(); - state.rmsnorm_weights_creates.clear(); - state.ple_pack_calls.clear(); - state.ple_create_calls.clear(); - state.tensor_creates.clear(); - state.host_view_creates.clear(); - state.tensor_windows.clear(); - state.tensor_writes.clear(); - state.tensor_reads.clear(); - state.dispatches.clear(); - state.pad_multiple = 64; - state.matmul_k_delta = 0; - state.matmul_n_delta = 0; - state.pad_answers_seen.clear(); - state.release_order.clear(); - state.fail_dispatch_at.reset(); - state.read_returns_nonfinite = false; - state.nonfinite_read_index.reset(); - last_error.clear(); -} - -flm::corelib::CorelibApi::Resolver Resolver() { - return [](std::string_view name) -> void* { return FunctionFor(name); }; -} - -ryzenai_corelib_stream_ptr MakeStreamForTest() { - return static_cast(NewObject("stream")); -} - -} // namespace fake_corelib diff --git a/src/test/gemma4_rai/fake_corelib.hpp b/src/test/gemma4_rai/fake_corelib.hpp deleted file mode 100644 index eb36b6d58..000000000 --- a/src/test/gemma4_rai/fake_corelib.hpp +++ /dev/null @@ -1,920 +0,0 @@ -#pragma once - -// A corelib test double scoped to what src/test/gemma4_rai's targets need. -// -// This is a SEPARATE fake from src/test/phi4_rai/fake_corelib.{hpp,cpp} -- -// deliberately, not out of preference for "one fake per model". Two concrete -// gaps rule out reusing or extending the Phi-4 fake for Gemma 4's shape-plan -// tests, and both are structural, not stylistic: -// -// 1. Phi-4's `RowsPadCall` records only `{helper, m, k, n, group_size}` -- -// it does not carry the ssmlp weights descriptor's `activation` or -// `post_feedforward_layernorm` fields at all. Gemma 4's shape plan must -// pin BOTH (see gemma4_rai_shape_plan.cpp), so a call record that drops -// them cannot express the assertion. -// 2. None of Phi-4's `MatmulPadCall`/`RowsPadCall`/`MhaPadCall` records the -// `stream` argument, even though the fake receives it -- there is -// nothing for "the stream was passed to every pad helper" to compare -// against. -// -// Editing src/test/phi4_rai/fake_corelib.{hpp,cpp} to add these is out of -// scope here: that fake is owned by another task and its own tests already -// pin its current shape. This file is therefore a full, independent -// implementation of the same `FLM_CORELIB_FUNCTIONS` ABI surface -// (corelib_api.hpp requires every entry to resolve). -// -// WHAT THIS FAKE VALIDATES, AND WHAT IT CANNOT. -// -// Round 1 of review found this file recording calls and validating almost -// nothing while backing the same ABI as the Phi-4 fake, which validates a -// great deal -- and named that gap as the shape that let two Criticals -// through a fully green suite on the Phi-4 side. The checks below were ported -// from that fake and adapted. They are: -// -// - `create_tensor_window` REJECTS a window that does not fit inside the -// allocation it is carved from, measured against the ROOT allocation that -// every window in a chain shares. -// - `tensor_write` / `tensor_read` REJECT a (count, offset) pair that runs -// past the tensor's OWN declared extent -- which is what corelib.h says -// they are bounded by, not the page-rounded allocation behind it. -// - a dispatch's row-bearing operands must AGREE with each other, and each -// must carry a row count the matching pad helper is KNOWN to have -// produced (`State::pad_answers_seen`, via `OnGrid`). -// - `ryzenai_corelib_ssmlp_bf16` REJECTS two operands that are windows over -// the SAME backing allocation. Comparison is on backing STORAGE identity -// (`DispatchRecord::operand_storage`), never on the `void*` handle, since -// `create_tensor_window` mints a fresh handle every call and two windows -// onto one allocation therefore compare unequal by handle. -// - `matmul_pad_shape` can be made to perturb `k`/`n` -// (`State::matmul_k_delta` / `matmul_n_delta`), so the shape plan's -// "helper changed padded K/N" rejection is reachable from a test instead -// of being dead code. -// - THE PACKED BYTES NAME WHAT WAS PACKED INTO THEM. Every weights object -// carries a four-byte `identity_tag` taken from its own components; -// `weights_copy_data` stamps it at offset 0 of the blob it hands out, so -// the weight cache's file contains it, and each -// `..._weights_create_from_file` reads it back out of the slice it was -// given. Without that the CACHE path carries no identity at all -- the -// descriptor says "some weight of this shape", corelib returns no -// components -- and `BindFromCache`'s assignment BY SLOT INDEX has -// nothing watching it. It is also the more faithful model: real corelib's -// packed bytes are a function of the weight, so two distinct tensors -// never produce one blob, which the old descriptor-only size was -// pretending otherwise. -// - a `flat_mha` / `ssmlp` dispatch RETAINS THE DESCRIPTOR it ran with -// (`DispatchRecord::mha_desc` / `ssmlp_desc`). Nothing here judges it -- -// see the `kv_shared` bullet below -- but retaining it is what makes an -// assertion about the value that actually crossed the ABI writable at -// all, which is what the reference driver's own per-layer `kv_shared` -// guard needs (`gemma4_driver.py:1670-1690`). -// -// AND HERE IS WHAT IT STILL CANNOT CATCH, stated plainly because this file -// has twice taken a Critical for a comment that read as broader coverage than -// the code delivered: -// -// - `pad_answers_seen` is filled ONLY by this engine's own calls into these -// same pad helpers, and `Gemma4ShapePlan::Build()` enumerates EVERY row -// bucket at construction. So for a helper the plan actually queried (and -// only for such a helper -- see the `OnGrid` bullet further down), that -// helper's answer set already holds every bucket's answer by the time any -// dispatch runs. A dispatch -// whose operands are UNIFORMLY sized to the wrong bucket -- decode handed -// bucket 128's rows instead of bucket 1's -- passes both `RowsAgree` and -// `OnGrid` and would fail on hardware. This is not a gap in the checks; -// it is a gap the fake CANNOT close, because the live row count a -// dispatch is meant to use is exactly the argument corelib 0.5.0 removed -// from the dispatch ABI. Hardware is the only real check. -// - NOTHING HERE CHECKS THE DESCRIPTOR AGAINST THE SHIPPED KERNEL SET. -// `flat_mha_pad_rows` records the descriptor and always succeeds. A -// `scale` or `window` that no artifact ships for -- the exact Critical -// round 1 found -- is invisible to this fake and is caught only by a test -// that reads the recorded field itself (see -// TestAttentionScaleIsExactlyOneEverywhere) or by real hardware. -// - `kv_shared` likewise. The fake cannot TELL an owning dispatch given a -// kvshare descriptor from a correct one; neither can corelib -- both -// buffers exist, both are the right shape, and the kernel that runs is -// simply a different one. What the fake now does is RECORD the value that -// crossed the ABI (`DispatchRecord::mha_desc`), which is a strictly -// weaker thing and is the whole of the improvement: it moves the check -// from impossible to writable, and leaves writing it to the caller. The -// reference driver writes exactly that check host-side, against -// `cfg.kv_layers`, and raises at load rather than at the first prompt -// (`gemma4_driver.py:1670-1690`); C7 should port it. Until it does, the -// only thing pinning `kv_shared` is the descriptor-level assertions in -// test_gemma4_shape_plan.cpp. -// - THE IDENTITY ON THE CACHE LEG IS FOUR BYTES, NOT THE WHOLE COMPONENT. -// `identity_tag` is one block tag, or one BF16 gamma element, or (for -// ssmlp) the gate stream's tag alone -- enough to say WHICH weight, never -// enough to say the whole of it was right. Two norms agreeing in their -// first element and differing later are one weight to a cache-bound -// engine here; so are two ssmlp blobs whose gates agree and whose ups do -// not. Neither can occur against this fixture (every F32 tensor is -// constant-filled, every Q8_0 tensor's tag is unique) and the pack leg -// still checks the full gamma and all three block streams. On a real -// packed blob the content IS the weight, so the gap is the fake's, not -// the design's. -// - The row grid here is one flat "round up to 64, except M == 1" rule for -// every op. Real corelib ships a different grid per op and per shape. A -// bug that depends on two ops' grids genuinely differing cannot be seen -// here. -// - `OnGrid` REJECTS a matmul output width nothing ever queried a pad -// answer for. That is only sound while the plan queries every width the -// engine will dispatch: a width the plan misses is a FALSE REJECTION of a -// correct dispatch, not a missing check. It has already happened once -- -// lm_head's `hidden -> vocab` was unqueried and C7's first dispatch was -// refused -- which is why `Gemma4ShapePlan::Build()` now plans lm_head. -// Read OnGrid's own comment in fake_corelib.cpp before adding a dispatch -// at a width the plan does not query. -// - The write/read bound above holds only for objects that HAVE a declared -// shape. `WithinExtent` returns true when the shape is empty, and every -// object minted by `matmul_weights_create_*`, `ssmlp_weights_create_*`, -// `create_stream` and `MakeStreamForTest` has an empty shape -- so a -// `tensor_write`/`tensor_read` aimed at one of those is UNBOUNDED here. -// DEFERRED, not overlooked: corelib's weights and stream objects are not -// tensors and no correct engine writes through them, so tightening it -// would be guarding a path that only misuse reaches. Stated because the -// bullet above reads, on its own, as an unqualified guarantee. -// - `create_tensor_window` applies its fit check only when the parent has -// storage. A window carved from a NULL parent has none, skips the check, -// and then also slips past `AllDistinctStorage`, which skips -// storage-less operands. Misuse rather than a live path; deferred for the -// same reason. -// - `MatmulHelperFor` keys on the output width N ALONE, so two matmuls with -// the same N and different K share one grid entry. GEMMA 4 HAS SUCH A -// PAIR, so this is not hypothetical: the output projection is -// `q_heads*head_dim -> hidden` on a sliding layer and -// `q_heads*global_head_dim -> hidden` on a full one -- 2048->1536 and -// 4096->1536 on E2B -- and both land on `"matmul-1536"`. Nothing else -// collides on either shipped row (E2B queries N in {2048, 4096, 256, 512, -// 1536, 262144}; E4B in {2048, 4096, 512, 1024, 2560, 262144}). -// DEFERRED: with one shared row rule here the two entries would hold the -// same answers anyway, so keying on (K, N) would buy nothing today. It -// would start to matter the moment this fake models per-shape grids. -// - `release_order` RECORDS AN ORDER; IT DOES NOT MODEL THE FAULT. corelib -// access-violates somewhere inside `ryzenai_corelib_object_release` when -// an object outlives its Stream, and this fake's `object_release` is a -// `delete` that cannot and does not reproduce that. So a green -// release-order test says "this engine releases in the order that was -// measured not to fault", NEVER "this engine does not fault". Nothing in -// this project has torn a Gemma 4 engine down against real corelib. -// - `fail_dispatch_at` INJECTS a status; it does not model a reason. The -// engine's own comment says a mid-chain rejection against real corelib -// is PROCESS-FATAL -- the unwind destructs runs whose commands are still -// in flight and XRT aborts uncatchably. Here the unwind is clean, so a -// test can observe `poisoned()` at all. What the test proves is that the -// FLAG is set on that path, never that the path is survivable. -// The same caveat covers `fail_dispatch_at = 0`, which C9's fix round -// uses for the other half of the definition: it shows that the engine -// does not POISON when its first dispatch is refused, and says nothing -// about whether real corelib can refuse one without having enqueued it. -// - `nonfinite_read_index` NARROWS the NaN injection to one read so the -// engine's SECOND `RequireFinite` call site is reachable at all. It -// chooses which guard fires; it does not make a NaN at that point -// plausible, and no test may read anything into WHY the logits would be -// non-finite while the hidden state was not. -// -// One further, load-bearing absence: there is no `ple_pad_calls` here, and -// there never will be. The pinned corelib.h has no -// `ryzenai_corelib_ple_bf16_pad_rows` (or any other ple padding entry point) -// at all -- ple_bf16's own doc comment says M "comes from `x`'s own extent -// and SELECTS THE KERNEL", implicitly, with no helper to ask. So the shape -// plan computes `ple_rows` (and, by the same reasoning, `rmsnorm_rows`) from -// a local, hardcoded bucket table taken verbatim from corelib.h's own -// ple_bf16 comment, never from a corelib call -- there is nothing here to -// fake. -// -// THE PLE AND RMSNORM *WEIGHTS* ENTRY POINTS ARE A DIFFERENT MATTER AND ARE -// FAKED FOR REAL. An earlier version of this paragraph said -// `FLM_CORELIB_FUNCTIONS` "has no ple or rmsnorm entries either", which was -// true when C6 wrote it and is not any more: C7 added -// `ple_bf16_weights_pack`/`_create`/`ple_bf16` and -// `rmsnorm_bf16_weights_create_reference`/`_from_file`/`rmsnorm_bf16`, -// because Gemma 4 dispatches both operators every layer. What this fake does -// with them: -// -// - `ple_bf16_weights_pack` HONOURS THE TWO-CALL PROTOCOL. `out == NULL` -// reports the size and dereferences none of the four arrays, and records -// (`PlePackCall::read_any_array`) that it did not, so "the engine sized -// its buffer before it held the data" is a checkable claim rather than a -// comment. It also refuses `group_size != 32`, which is the only value -// mladfple ships. -// - it RETAINS `post_norm` and `next_norm` in full. `next_norm` is the -// NEXT layer's attention-norm gamma and the off-by-one there still -// generates text, so it has to be assertable per layer. The layer a pack -// call belongs to is read off `desc.layer_scale`, which the fixture -// makes unique per layer. -// - `rmsnorm_bf16_weights_create_reference` retains the whole BF16 gamma, -// so "V's norm is a vector of exactly 1.0" is checkable rather than -// assumed. -// - `ple_bf16` and `rmsnorm_bf16` ARE RECORDED AS DISPATCHES, which until -// Task R3 they were not: only matmul / ssmlp / mha minted a -// `DispatchRecord` and everything else fell through the generic success -// branch. A Gemma 4 layer is SEVEN dispatches on a sharing layer and -// `10 + kv_heads` on an owning one -- 11 on E2B, 12 on E4B -- of which -// these two ops are 3 of 7 and 5 of 11 (6 of 12) respectively, so C8's -// layer loop would have been measured by an instrument blind to roughly -// HALF of it. (An earlier version of this comment said "ten and six": -// those are the WEIGHT-OBJECT counts from gemma4_rai.hpp:40-47, correct -// for weights and wrong for dispatches. Counted from gemma4_driver.py -// :2254-2336.) Each -// carries its operands in the header's own order, their backing STORAGE -// identity, the weights object it ran with, and its position in -// `State::dispatches`, which is ABI call order. -// - `ple_bf16` enforces what its section of corelib.h states: `out` and -// `norm_out` distinct from each other and from `x` (on storage, not on -// handles), and a row count the ELF set actually ships -- see OnPleGrid, -// which is stricter than corelib on purpose. -// - `rmsnorm_bf16` enforces NEITHER of those, and the absences are real -// rather than pending: in-place is explicitly legal for small M, and -// `rmsnorm_bf16_pad_rows` is an identity function that is not even bound -// in FLM_CORELIB_FUNCTIONS, so there is no grid to check against. See the -// dispatch branch in fake_corelib.cpp for both, written out. -// - `weights_copy_data` now reports a real (small, synthetic) size instead -// of 0. It used to report 0 unconditionally, which silently disabled the -// weight cache's whole write path. - -#include "rai/corelib_api.hpp" - -#include -#include -#include -#include -#include -#include -#include - -namespace fake_corelib { - -/// \brief one `ryzenai_corelib_matmul_bf16_pad_shape` call -/// \note `m`, `k` and `n` are recorded as they were PASSED IN, before any -/// padding or `matmul_*_delta` perturbation, so a test can classify a -/// call by its logical operand shape. -struct MatmulPadCall { - ryzenai_corelib_stream_ptr stream; - std::int64_t m; - std::int64_t k; - std::int64_t n; - std::uint32_t group_size; -}; - -/// \brief one `ryzenai_corelib_ssmlp_bf16_pad_rows` call -/// -/// Carries the WHOLE descriptor, not a projection of it -- `activation` and -/// `post_feedforward_layernorm` are exactly the fields Gemma 4's shape plan -/// must get right and Phi-4's fake never recorded (see this header's own -/// top comment). -struct RowsPadCall { - ryzenai_corelib_stream_ptr stream; - std::int64_t m; - ryzenai_corelib_ssmlp_bf16_weights_desc desc; -}; - -/// \brief one `ryzenai_corelib_flat_mha_bf16_pad_rows` call -/// -/// Carries the WHOLE descriptor for the same reason: `scale`, `kv_shared` -/// and `rope_dim` change WHICH KERNEL RUNS rather than any shape -/// (corelib.h:1810-1812), so a fake that dropped them would make the only -/// assertions capable of catching a wrong one impossible to write. -struct MhaPadCall { - ryzenai_corelib_stream_ptr stream; - std::int64_t m; - ryzenai_corelib_flat_mha_bf16_desc desc; -}; - -/// \brief one `..._enum_kernels` call, and the stream it was asked about -struct EnumKernelsCall { - std::string entry; - ryzenai_corelib_stream_ptr stream; -}; - -/// \brief a matmul kernel an explicit grid reports -struct ShippedMatmul { - std::int64_t m, k, n, group; -}; - -/// \brief a fused-MLP kernel an explicit grid reports -struct ShippedSsMlp { - std::string family; - std::int64_t m, k, n, group; -}; - -/// \brief an attention kernel an explicit grid reports -/// \note `window` is the window the kernel is KEYED on, which corelib only -/// writes when it is shorter than the rows the kernel runs. -struct ShippedMha { - std::int64_t heads, kv_heads, m, head_size, max_seq, window; - bool kv_shared, scale_one; -}; - -/// \brief one paged attention kernel `paged_mha_enum_kernels` reports -struct ShippedPaged { - std::int64_t heads, kv_heads, m, head_size, page_size; -}; - -/// \brief one `ryzenai_corelib_matmul_bf16_weights_create_*` call -/// \note The WHOLE descriptor, because `group_size` is the field C7 must get -/// right and it is the one nothing else in this codebase can observe: -/// `group`, `head_group` and `ple_group` are all 32 on both shipped -/// Gemma 4 rows, so a weight requantized at the wrong one is a wrong -/// weight that every real config agrees with. A test forces them apart -/// (`head_group = 64`) and reads this back. -struct MatmulWeightsCreateCall { - ryzenai_corelib_matmul_bf16_weights_desc desc; - /// \brief true for the `..._from_file` leg, i.e. a cache bind rather than - /// a pack -- which is how "the second load hit the cache" is told - /// apart from "the second load packed again" - bool from_file; - /// \brief the four-byte tag at the head of the block stream, 0 on the - /// `..._from_file` leg - /// \note WHICH QUANTIZED TENSOR WENT IN. The descriptor answers only - /// "some tensor of this shape at this group size", and Gemma 4 has - /// 35 (E2B) layers' worth of identically-shaped `attn_q` weights. - /// See gemma4_gguf_fixture.hpp's BlockSignatureAt. - std::uint32_t blocks_tag; - /// \brief the weights object this call MINTED - /// \note The bridge between "which bytes went in" and "which slot the - /// finished object ended up in" -- the two are different questions - /// and only the first was answerable before. See - /// DispatchRecord::weights and gemma4_rai's WeightPlacementsForTest. - void* object; -}; - -/// \brief one `ryzenai_corelib_ssmlp_bf16_weights_create_*` call -/// \note Carries the FIRST ELEMENT of each norm, not the whole vector. That -/// is enough only because the fixture fills every F32 tensor with a -/// constant signature derived from its own name -/// (gemma4_gguf_fixture.hpp's `NormSignature`), so one element -/// identifies WHICH tensor was passed. Against a real GGUF it would -/// identify nothing, and this record would have to grow. -struct SsMlpWeightsCreateCall { - ryzenai_corelib_ssmlp_bf16_weights_desc desc; - bool from_file; - bool has_components; - std::uint16_t epsilon_bf16; - std::uint16_t norm0_bf16; - std::uint16_t norm1_bf16; - /// \brief the four-byte tag at the head of each of the three block - /// streams, 0 on the `..._from_file` leg - /// - /// THE PORT ASSIGNMENT, WHICH NOTHING COULD SEE BEFORE. A gated MLP puts - /// `gate` through the activation and `up` not, so exchanging the two - /// computes `GELU(x @ up) * (x @ gate)` -- a different function, finite, - /// plausible and wrong. corelib takes all three as bare `const void*` - /// block streams of the same quant type, so nothing in the ABI, the - /// descriptor or the shapes distinguishes them (`gate` and `up` are the - /// SAME shape on every Gemma 4 layer). The fixture's per-tensor block - /// tag is the only discriminator there is. - std::uint32_t gate_tag; - std::uint32_t up_tag; - std::uint32_t down_tag; - /// \brief the weights object this call minted -- see - /// MatmulWeightsCreateCall::object - void* object; -}; - -/// \brief one `ryzenai_corelib_rmsnorm_bf16_weights_create_*` call -struct RmsNormWeightsCreateCall { - /// \brief the PDI handed to the reference packer, or -1 on `..._from_file` - int prefill_pdi; - ryzenai_corelib_rmsnorm_bf16_weights_desc desc; - bool from_file; - /// \brief the BF16 gamma, `desc.k` entries, EMPTY on the from_file leg - /// \note Retained in full rather than sampled: V's norm is a vector of - /// exactly 1.0 and the cheap wrong implementation is to skip it, so - /// a test has to be able to say "every element is one" rather than - /// "the first one is". - std::vector scale; - /// \brief the gamma's first element, ON BOTH LEGS - /// - /// WHICH NORM THIS IS, ON THE ONE PATH WHERE `scale` CANNOT SAY. corelib - /// hands no gamma back through `..._weights_create_from_file` -- it is - /// given a slice of an opaque packed blob -- so a cache-bound engine - /// carries no gamma anywhere and every verifier keyed on `scale` is blind - /// to it. That matters because `BindFromCache` assigns handles by walking - /// the slot table IN INDEX ORDER, which is a second, entirely independent - /// way for a weight to land in the wrong member, and one layer's Q-norm - /// and K-norm have IDENTICAL descriptors and therefore identical packed - /// lengths: exchanging them binds cleanly. - /// - /// What closes it is that this fake's packed bytes now CARRY the identity - /// of what was packed (see fake_corelib.cpp's `identity_tag`), so the - /// from_file leg recovers this from the cache file itself. That is also - /// the more faithful model: real corelib's packed bytes are a function of - /// the weight, and two distinct gammas do not produce one blob. - std::uint16_t scale_tag; - /// \brief the weights object this call minted -- see - /// MatmulWeightsCreateCall::object - /// \note Load-bearing for THIS record in particular. Gemma 4's Q-norm and - /// K-norm on one layer have identical descriptors (same `k`, same - /// epsilon) and differ only in their gamma, so "the object packed - /// from `attn_q_norm` is the one that reached the `q_norm` slot" - /// cannot be said without an object identity to join on. - void* object; -}; - -/// \brief one `ryzenai_corelib_ple_bf16_weights_pack` call -/// -/// BOTH LEGS OF THE TWO-CALL PROTOCOL LAND HERE, told apart by `sizing_leg`. -/// The sizing leg (`out == nullptr`) must report a size and READ NONE OF THE -/// ARRAYS; `read_any_array` records whether this fake dereferenced them, so -/// a test can prove the engine really made the sizing call first rather than -/// calling the full packer twice. -struct PlePackCall { - ryzenai_corelib_ple_bf16_weights_desc desc; - bool sizing_leg; - bool read_any_array; - std::size_t reported_size; - /// \brief the first element of `gate` and of `proj`, 0 on the sizing leg - /// \note WHICH TENSOR WENT INTO WHICH PORT. `gate` is [k, n] and `proj` - /// is [n, k]; on this model those are 1536x256 and 256x1536, the - /// same element count, so a swap is not a size error -- corelib.h - /// says passing one the other way round "packs silently and - /// produces noise". The fixture's per-tensor signature is what - /// makes the swap visible here. - /// \note IDENTITY, NOT ORIENTATION. Element 0 is `M[0][0]` under either - /// orientation, so these two say WHICH tensor arrived and can say - /// nothing about whether it arrived transposed. `gate_at_*` and - /// `proj_at_*` below are the orientation probes. - float gate_first; - float proj_first; - /// \brief two OFF-DIAGONAL probes of each matrix, 0 on the sizing leg - /// - /// THE PAIR THAT SWAPS UNDER A TRANSPOSE, WHICH IS THE WHOLE REASON THEY - /// ARE HERE. Flat index 0 and flat index `k*n - 1` are both fixed points - /// of a transpose -- first and last element of the buffer either way - /// round -- so a record that probes only those is blind by construction, - /// however the fixture is filled. Index 1 is `(row 0, col 1)` and index - /// `n` (for `gate`, whose required layout is `[k, n]`) is - /// `(row 1, col 0)`: transposing exchanges exactly these two. `proj`'s - /// required layout is `[n, k]`, so its second probe is at index `k`. - /// - /// Reading them is only informative because `gemma4_gguf_fixture.hpp` - /// gives 2-D F32 tensors a POSITIONAL fill -- against a constant-filled - /// matrix every index reports the same number. - float gate_at_1; - float gate_at_n; - float proj_at_1; - float proj_at_k; - /// \brief the inner RMSNorm's gamma, `desc.k` floats -- empty on the - /// sizing leg - std::vector post_norm; - /// \brief THE NEXT LAYER'S attention-norm gamma, `desc.k` floats -- empty - /// on the sizing leg - /// \note This is the off-by-one the header says "is not a crash -- every - /// layer is then normalized by its neighbour's gamma, which still - /// generates text". Retained per call so it can be asserted for - /// EVERY layer rather than spot-checked, keyed by - /// `desc.layer_scale` (which the fixture makes unique per layer). - std::vector next_norm; -}; - -/// \brief one `ryzenai_corelib_ple_bf16_weights_create` call -struct PleCreateCall { - ryzenai_corelib_ple_bf16_weights_desc desc; - std::size_t packed_size; - /// \brief false when the blob's length was not exactly what this - /// descriptor packs to -- corelib.h: "a truncated blob is still a - /// plausible one", so this fake refuses the same way - bool accepted; - /// \brief the weights object this call minted, null when refused -- see - /// MatmulWeightsCreateCall::object - void* object; -}; - -/// \brief one `ryzenai_corelib_create_device_tensor` call -struct TensorCreateRecord { - ryzenai_corelib_data_type data_type; - std::vector shape; - void* object; -}; - -/// \brief one `ryzenai_corelib_create_tensor_window` call that SUCCEEDED -/// \note A rejected window records nothing -- there is no object to record. -struct TensorWindowRecord { - void* parent; - std::vector shape; - std::size_t offset; - void* object; - /// \brief the strides the window was created with, empty for dense - std::vector strides; -}; - -/// \brief one `ryzenai_corelib_tensor_write` call -struct TensorWriteRecord { - void* tensor; - ryzenai_corelib_data_type source_type; - std::size_t count; - std::size_t offset; - bool accepted; -}; - -/// \brief one `ryzenai_corelib_tensor_read` call -/// \note WHICH ROW WAS READ, which is the whole content of the quietest bug a -/// forward pass has. A pass ends by reading the LAST live row of -/// `hidden`, at element offset `(rows - 1) * hidden`; reading row 0 -/// instead returns the first prompt token's hidden state -- finite, -/// correctly shaped, correctly bounded, and the wrong position's logits. -/// Nothing downstream can tell. -/// \note Refused reads are recorded too (`accepted == false`): "read past the -/// end" and "did not read at all" are different bugs. -struct TensorReadRecord { - void* tensor; - ryzenai_corelib_data_type destination_type; - std::size_t count; - std::size_t offset; - bool accepted; -}; - -/// \brief one accepted dispatch of matmul / ssmlp / rmsnorm / ple / flat_mha -/// \note A REFUSED DISPATCH RECORDS NOTHING. Every per-op check below returns -/// before the record is pushed, so `State::dispatches` is the list of -/// calls that would have reached the NPU. -struct DispatchRecord { - /// \brief "matmul", "ssmlp", "rmsnorm", "ple" or "mha" - std::string kind; - void* stream; - std::int64_t rows; - /// \brief every operand of this dispatch corelib.h ties to a row count, - /// in a fixed per-kind order: matmul is {input, output}; ssmlp is - /// {input, residual, skip_sum, normalized}; rmsnorm is - /// {input, output}; ple is {x, ple, out, norm_out}; mha is - /// {q, k, out} - /// \note THE ORDER IS PART OF THE CONTRACT, because it is how a caller - /// says "the thing that was normalized is the thing the block - /// wrote" rather than "two tensors were involved". It follows each - /// op's own argument order in corelib.h, minus the weights object - /// (which is `weights`) and minus operands with unrelated extent - /// rules (mha's cos/sin/caches). - std::vector operands; - /// \brief the underlying allocation identity behind each of `operands`, - /// in the same order - /// \note An aliasing check MUST compare this, not `operands` itself. - /// `create_tensor_window` mints a fresh object every call, so two - /// DIFFERENT window objects can -- and, when a bug routes two - /// "must be distinct" operands through separate windows onto the - /// SAME backing, will -- point at the SAME storage. Entries here - /// would then be equal even though the corresponding `operands` - /// entries are not. See fake_corelib.cpp's StorageOf. - std::vector operand_storage; - /// \brief the DECLARED SHAPE of each of `operands`, in the same order - /// - /// THE COLUMN EXTENT, WHICH NOTHING HERE RETAINED AND WHICH IS HALF OF - /// EVERY GEOMETRY. `rows` is `operands.front()`'s leading extent and - /// `RowsAgree`/`OnGrid` compare leading extents only, so until this field - /// existed a `DispatchRecord` said nothing whatever about how WIDE its - /// operands were. Gemma 4 runs TWO geometries -- sliding layers at - /// `head_dim` and full-attention layers at `global_head_dim`, 256 and 512 - /// on both shipped rows -- over ONE set of allocations sized for the - /// wider, windowed down per layer on the column axis. So a layer that - /// binds the other geometry's windows produces operands that are the same - /// tensors, at the same offsets, over the same storage, with the same row - /// counts, and the wrong width. Every comparison this record could make - /// before still passed: a mutation that bound the FULL geometry on every - /// layer left the whole suite green. - /// - /// \note PURE RETENTION, like `mha_cos` and `weights`. Nothing in this - /// fake judges these shapes -- no width check was added to any - /// dispatch branch -- so no dispatch that was accepted before is - /// refused now, and the row checks keep the meaning they had. What - /// changes is only that "this layer ran at ITS OWN geometry's - /// widths" becomes a question a caller can ask. - /// \note THE WHOLE SHAPE, NOT THE COLUMN COUNT. A rank-3 operand (the V - /// staging destination at more than one KV head) has no "column" - /// -- recording `shape[1]` for it would silently mean the middle - /// extent -- so the shape is carried as it was declared and the - /// reader says which extent it means. - std::vector> operand_shapes; - /// \brief the flat_mha descriptor this dispatch was actually handed - /// \note THE DESCRIPTOR, NOT THE PLAN'S COPY OF IT. Holding four - /// immutable descriptors makes the wrong one impossible to - /// CONSTRUCT; it does nothing about the right one being built and - /// the wrong one being HANDED OVER. The reference driver closes - /// that gap host-side by reading the field back per layer against - /// `cfg.kv_layers` and raising at load - /// (`gemma4_driver.py:1670-1690`); a port of that guard, and any - /// C7/C8 assertion of the same shape, needs the value that crossed - /// the ABI, which is this. - /// \note Meaningless unless `has_mha_desc`. corelib passes the descriptor - /// by pointer and a null one is legal to attempt, so "all zeroes" - /// and "absent" are different states and a reader must be able to - /// tell them apart. - ryzenai_corelib_flat_mha_bf16_desc mha_desc{}; - bool has_mha_desc{false}; - /// \brief `flat_mha`'s four remaining tensor operands: the rotary pair it - /// roped with and the two caches it attended over - /// - /// NOT IN `operands`, AND THAT IS DELIBERATE RATHER THAN AN OMISSION. - /// `operands` is the row-bearing set -- the tensors `RowsAgree` and - /// `OnGrid` are entitled to compare -- and these four have their own, - /// different extent rules: the caches are rank 3 and pinned to - /// `desc.max_seq`, and the rotary tables are FP32 `[positions, - /// rope_dim/2]`. Putting them in `operands` would either weaken those - /// checks or make them lie. - /// - /// RECORDED BECAUSE BOTH ARE SILENT WHEN WRONG, in exactly the way - /// `DispatchRecord::weights` is. A layer handed the other regime's rotary - /// pair ropes at the wrong base and "still converges to fluent text"; a - /// layer handed the wrong cache attends over another layer's keys and - /// values and nothing errors. Neither this fake nor corelib can tell a - /// correct binding from an incorrect one -- both tensors are real, both - /// are the right shape -- so the only possible move is to record the - /// identity that crossed the ABI and let the caller join it to the tensor - /// it created for that role (`State::tensor_creates`, or - /// `State::tensor_windows` for a window onto one). - /// - /// \note Added by Task C8, which is also the first task to dispatch - /// attention. That is the arrangement R3 warned about -- the layer - /// loop extending the instrument it is measured by -- and it is - /// confined to RETAINING a handle: nothing here judges the value, - /// and no check anywhere else was relaxed to accommodate it. - void* mha_cos{nullptr}; - void* mha_sin{nullptr}; - void* mha_k_cache{nullptr}; - void* mha_v_cache{nullptr}; - /// \brief the absolute position `flat_mha` was told these rows start at, - /// or -1 for a dispatch that is not attention - /// - /// "Q'S SHAPE IS A POSITION RANGE. Its rows are the rows this call ropes, - /// and `position` is the absolute index of the first of them -- together - /// they span [position, position + M), which is what the rotary and the - /// causal mask need." It is therefore the one scalar in the pass that - /// makes a decode step different from a prefill, and an engine that always - /// passed 0 would rope every token as if it were the first and mask every - /// step to a one-token context. That is finite, plausible, fluent and - /// wrong -- and until this field existed no test could see it, which was - /// demonstrated by a mutation that passed a literal 0 here and left the - /// whole suite green. - std::int64_t mha_position{-1}; - /// \brief the ssmlp descriptor this dispatch's WEIGHTS were created with - /// \note `ryzenai_corelib_ssmlp_bf16` takes a packed weights object, not a - /// descriptor (corelib.h:1384-1388) -- there is no descriptor at - /// the dispatch to record. What is carried across instead is the - /// descriptor handed to `ssmlp_bf16_weights_create_*`, which is - /// where `activation` and `post_feedforward_layernorm` live and so - /// is what a C8 assertion about this block actually wants. - /// \note Meaningless unless `has_ssmlp_desc` -- false when the dispatch - /// carried no weights object, or one this fake did not mint. - ryzenai_corelib_ssmlp_bf16_weights_desc ssmlp_desc{}; - bool has_ssmlp_desc{false}; - /// \brief the WEIGHTS OBJECT this dispatch was handed, or null for a kind - /// that takes none - /// - /// "THE RIGHT DESCRIPTOR BUILT AND THE WRONG OBJECT HANDED OVER" is a - /// failure mode neither this fake nor corelib can detect from inside a - /// dispatch -- both objects are real, both were packed by the same entry - /// point, and the kernel simply multiplies by the wrong numbers. What is - /// recorded here is the identity that crossed the ABI, which is what - /// makes the assertion writable by the caller: joined against - /// `MatmulWeightsCreateCall::object` (or the ssmlp / rmsnorm / ple - /// equivalents) and against `gemma4_rai`'s WeightPlacementsForTest, a - /// test can say "layer 7's attn_q dispatch ran with the object packed - /// from blk.7.attn_q.weight". - /// - /// Null for `flat_mha`, which takes no weights object at all -- its - /// operands are q/k/v, the caches and the rope tables - /// (corelib.h:1384-1388 and the flat_mha section). That is an absence in - /// the ABI, not an unrecorded field. - /// - /// SET FOR `rmsnorm` AND `ple` TOO, and that is what makes them - /// identifiable at all: neither dispatch carries a descriptor, so "which - /// layer's norm was this?" is answered only by joining this handle to - /// `RmsNormWeightsCreateCall::object` (whose `scale` or `scale_tag` names - /// the gamma) or to `PleCreateCall::object` (whose `desc.layer_scale` - /// names the layer). No descriptor is duplicated onto this record for - /// them: a second copy would be a second thing to keep in step with the - /// first, and the join is already grounded in the object that crossed the - /// ABI. - /// - /// \note C8, not C7, is where this becomes reachable from the engine: - /// nothing dispatches yet. It is landed now because C8's tests have - /// to be WRITTEN against it, and because the blind spot it closes - /// multiplies by 35 layers the moment the layer loop exists. - void* weights{nullptr}; -}; - -/// \brief one `ryzenai_corelib_stream_enable_kernel_bundle` call -struct KernelBundleRecord { - void* stream{}; - std::string model; - std::size_t enum_calls_before{}; -}; - -/// \note EVERY MEMBER IS WRITTEN UNDER ONE MUTEX inside the fake, because -/// the engine packs its weights across a pool of threads -/// (`kWeightCreateConcurrency`). A reader -- i.e. a test -- runs after -/// the pool has joined and takes no lock. Before C7 nothing here was -/// called concurrently and there was no mutex at all; adding one is -/// not a style change, it is what stops eight threads pushing into the -/// same vector. -/// \note NOTHING HERE IS IN A MEANINGFUL ORDER once the pool is involved. -/// Assertions must key on a field that identifies the weight (the -/// descriptor, a norm's signature value) rather than on an index -- -/// exactly as the engine assigns its results by slot rather than by -/// completion. -struct State { - /// \brief every `..._enum_kernels` call, in order - std::vector enum_calls; - std::vector kernel_bundles; - /// \brief report exactly the `shipped_*` kernels instead of "every M" - /// \note What makes a plan's kernel QUERY observable: the plan filters - /// the reported set itself, so the only thing a test can see is - /// whether the set it shipped covers what the plan asked for. - bool explicit_grid{false}; - std::vector shipped_matmul; - std::vector shipped_ssmlp; - std::vector shipped_mha; - /// \brief report exactly these paged kernels; empty reports "every M" - std::vector shipped_paged; - std::vector matmul_pad_calls; - std::vector rows_pad_calls; - std::vector mha_pad_calls; - std::vector matmul_weights_creates; - std::vector ssmlp_weights_creates; - std::vector rmsnorm_weights_creates; - std::vector ple_pack_calls; - std::vector ple_create_calls; - std::vector tensor_creates; - /// \brief every `ryzenai_corelib_create_host_view` that succeeded - std::vector host_view_creates; - std::vector tensor_windows; - std::vector tensor_writes; - std::vector tensor_reads; - /// \brief every accepted dispatch, IN ABI CALL ORDER - /// \note THE ONE VECTOR HERE WHOSE ORDER MEANS SOMETHING, and the - /// exception to the "nothing here is in a meaningful order" note - /// below. That note is about the weight CREATES, which run across - /// `kWeightCreateConcurrency` threads; a layer loop dispatches from - /// one thread, so push order under the state mutex is call order. - /// A Gemma 4 layer is SEVEN dispatches on a sharing layer and - /// `10 + kv_heads` on an owning one (11 on E2B, 12 on E4B; V's norm - /// runs once per KV head). A PERMUTATION of them is a different - /// model while a COUNT of them is not -- so C8 should assert the - /// sequence, not the multiset. - std::vector dispatches; - - /// \brief the row grid every pad helper here rounds through - /// \note M == 1 is never rounded: corelib's flat_mha doc says decode - /// reports 1 because the token kernel runs a single row. - std::int64_t pad_multiple{64}; - - /// \brief what `matmul_pad_shape` adds to the caller's `k` and `n` - /// \note Both are 0 by default, i.e. "the helper returns K and N - /// unchanged", which is what real corelib does for every shape this - /// model presents. A test sets one non-zero to reach - /// `MatmulRows`'s "helper changed padded K/N" rejection in - /// gemma4_rai_shape_plan.cpp, which is otherwise unreachable and - /// therefore untested. - std::int64_t matmul_k_delta{0}; - std::int64_t matmul_n_delta{0}; - - /// \brief every value a pad helper has actually returned, per helper name - /// \note This is what lets a dispatch be checked against its OWN op - /// family's pad answers instead of merely against its co-operands. - /// Helper names are `"matmul-"` (one per distinct output width, - /// so a Q-projection operand cannot be validated against the output - /// projection's grid), `"ssmlp"` and `"mha"`. Read the header's - /// "what it still cannot catch" list before treating a green - /// dispatch test as proof the extents were right. - std::unordered_map> pad_answers_seen; - - // ---- Task C9's three fields. See the header's "what this fake still - // cannot catch" list: two of them make a FAILURE PATH reachable at - // all, and the third records an order the ABI already had. None of - // them changes what any other assertion in the suite observes. - - /// \brief every `ryzenai_corelib_object_release`, IN ABI CALL ORDER, by - /// the kind string the object was created with - /// - /// THE SECOND VECTOR HERE WHOSE ORDER MEANS SOMETHING, and it is the - /// whole instrument behind C9. corelib access-violates when a packed - /// object is released AFTER the `Stream` it was planned against -- 7 of - /// 10 measured runs in the reference driver, up to 221 printed lines in - /// one run. An earlier version of this note added "and the process - /// exited 0 every time with nothing raised"; the source says nineteen of - /// TWENTY runs -- both teardown orders together -- exited 0 and the - /// twentieth exited 1, and does not break that down by order. Which is - /// still enough: at 19 of 20 an exit code asserts nothing, so a test - /// cannot wait for the fault and reads the ORDER instead. - /// Releases run on the destroying thread and nothing else releases - /// concurrently, so push order under the state mutex is call order. - /// - /// Kinds are the strings `NewObject` is called with: `"stream"`, - /// `"tensor"`, `"window"`, `"matmul_weights"`, `"rmsnorm_weights"`, - /// `"ple_weights"`, `"ssmlp_weights"` -- plus `"cleanup"`, which is NOT - /// an object release but `ryzenai_corelib_cleanup`, the runtime's own - /// teardown call. It is in this log so that "the runtime goes last of - /// all" is one question about one ABI sequence rather than two questions - /// about two instruments. - /// - /// \note NOT CLEARED BY ANYTHING BUT `Reset()`. A test that wants only - /// one teardown's worth takes `release_order.size()` as a mark - /// first and reads the slice appended after it. - std::vector release_order; - - /// \brief the 0-based index of the dispatch this fake REFUSES, if any - /// - /// WITHOUT IT THE ENGINE'S FAILURE PATH IS UNREACHABLE FROM A TEST and - /// `poisoned()` could only ever be observed false. Counted over accepted - /// dispatches in this run (`dispatches.size()`), so "fail the 20th" is a - /// MID-CHAIN rejection -- the case the engine's catch exists for -- - /// rather than a refusal before anything was submitted. - /// - /// \note The refused dispatch records NOTHING. It never happened as far - /// as `dispatches` is concerned, which keeps the index stable: - /// setting it to 20 refuses whatever would have been the 21st - /// accepted call, not a moving target. - /// \note This is an INJECTION, not a validation. It does not model any - /// reason corelib would reject a dispatch, and no test may read - /// anything into WHICH op it lands on beyond "it is not the first". - std::optional fail_dispatch_at; - - /// \brief make every accepted `tensor_read` return a quiet NaN - /// \note The ONLY way to reach `RequireFinite` in gemma4_rai.cpp, which - /// is the engine's guard against an all-NaN result reaching an - /// argmax that then answers `` for every step -- a CLI that - /// prints an empty answer under a plausible token count and exits - /// 0. Until C9 that guard had no test at all. - /// \note Set AFTER the engine is constructed. A load performs no - /// `tensor_read`, but nothing here depends on that staying true. - bool read_returns_nonfinite{false}; - - /// \brief which accepted read the NaN lands on: an index into - /// `tensor_reads`, or every read when unset - /// - /// WITHOUT IT THE INJECTION IS GLOBAL AND ONLY THE FIRST `RequireFinite` - /// CALL SITE IS EVER REACHED. gemma4_rai.cpp checks twice -- the layer - /// loop's last live hidden row, and then the head's logits -- and with - /// every read poisoned the first throws and the second is dead code as - /// far as any test is concerned. Review-C9's m-1 measured exactly that: - /// deleting the logits check left the suite green. - /// - /// \note NULLOPT IS THE OLD BEHAVIOUR, byte for byte: when this is unset - /// and `read_returns_nonfinite` is true, every accepted read - /// returns NaN as before. It narrows an existing injection and - /// adds no new one -- with `read_returns_nonfinite` false this - /// field does nothing at all. - /// \note COUNTED OVER `tensor_reads`, WHICH INCLUDES REFUSED READS, and - /// the index is that of the read in progress at the time the value - /// is chosen. A test should take `tensor_reads.size()` as a mark - /// rather than assuming a pass's first read is index 0. - /// \note Like `fail_dispatch_at`, this is an INJECTION and models no - /// reason. It proves which guard fires, never that a real device - /// would produce a NaN there. - std::optional nonfinite_read_index; -}; - -/// \brief what `tensor_read` puts in a caller's buffer -/// -/// A SYNTHETIC CONSTANT, AND IT MEANS NOTHING EXCEPT "NOT ZERO". This fake -/// moves no bytes -- there is no device, and the contents it would have to -/// invent are the whole of what a real corelib computes -- so a read cannot -/// return the right answer. It used to return zeros, and that turned out to -/// hide a specific, named bug: Gemma 4 softcaps its logits on the HOST, after -/// the read-back, and `tanh(0/30)*30` is 0, so an engine that dropped the -/// softcap entirely was indistinguishable from one that applied it. A -/// mutation confirmed that: removing the call left the whole suite green. -/// -/// Chosen well outside the cap's near-linear region, so a monotonic squash is -/// unmistakable: `tanh(45/30)*30` is about 27.15, not 45. -/// -/// \note This is the ONLY thing a test may conclude from a read's contents. -/// "The value the engine returned is this constant, transformed by the -/// host-side arithmetic the engine claims to apply" is a fair question; -/// anything about what the NPU would have computed is not. -inline constexpr float kSyntheticReadValue = 45.0f; - -State& GetState(); -void Reset(); -flm::corelib::CorelibApi::Resolver Resolver(); -ryzenai_corelib_stream_ptr MakeStreamForTest(); - -/// \brief record a shipped matmul row count the way a kernel enumeration would -ryzenai_corelib_status NoteMatmulPad(ryzenai_corelib_stream_ptr stream, - std::int64_t* rows, - std::int64_t* k, - std::int64_t* n, - std::int64_t group); -/// \brief record a shipped ssmlp row count -ryzenai_corelib_status NoteSsmlpRows(ryzenai_corelib_stream_ptr stream, - std::int64_t* rows, - const ryzenai_corelib_ssmlp_bf16_weights_desc* desc); -/// \brief record a shipped attention row count -ryzenai_corelib_status NoteMhaRows(ryzenai_corelib_stream_ptr stream, - std::int64_t* rows, - const ryzenai_corelib_flat_mha_bf16_desc* desc); - -/// \brief the byte count this fake's `ple_bf16_weights_pack` reports for a -/// descriptor, and the only length its `ple_bf16_weights_create` -/// accepts -/// \note A SMALL, DETERMINISTIC STAND-IN, not corelib's real packed size. -/// Real ple blobs are ~1 MB per layer and the weight-cache test writes -/// every packed weight to a file; using realistic sizes would make -/// that test write gigabytes to prove a bookkeeping property. What the -/// number has to be is (a) the same for the same descriptor and (b) -/// different for different ones, so that a cache entry written for one -/// weight cannot be bound to another. Exposed so a test can assert the -/// two-call protocol's reported size is what the second leg is then -/// given, rather than duplicating the formula. -std::size_t PlePackedSize(const ryzenai_corelib_ple_bf16_weights_desc& desc); - -} // namespace fake_corelib diff --git a/src/test/gemma4_rai/gemma4_gguf_fixture.hpp b/src/test/gemma4_rai/gemma4_gguf_fixture.hpp deleted file mode 100644 index f9af43f76..000000000 --- a/src/test/gemma4_rai/gemma4_gguf_fixture.hpp +++ /dev/null @@ -1,884 +0,0 @@ -#pragma once - -// A synthetic GGUF writer for Gemma 4's two shipped shapes (E2B and E4B). -// -// This is a second, independent implementation of the same GGUF v3 byte -// layout src/test/phi4_rai/gguf_fixture.hpp's Builder writes -- not an -// extraction of it. That header's Builder hardcodes Phi-4's own contract -// metadata in its constructor (general.architecture = "phi3", phi3.* keys, a -// 200064-entry tokenizer.ggml.tokens array...) and is owned by the agent -// working phi4_rai's corelib tests; reusing it here would mean either -// editing a file this task does not own or fighting its Phi-3-shaped -// defaults from the outside for every metadata key gemma4 needs instead. -// A from-scratch, model-agnostic writer -- generic over metadata and -// tensors, carrying no baked-in architecture -- is cleaner for what this -// fixture needs: two files whose overlap with Phi-4's fixture is the -// container format itself, nothing else. - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#if defined(_WIN32) -#include // _getpid -#else -#include // getpid -#endif - -namespace gemma4_fixture { - -namespace detail { - -/// \brief the process id, folded into every temp filename below so that two -/// instances of the *same* test binary running at once (`ctest -j`, -/// or a developer's scratch repro alongside the real suite) never -/// compute the same path. A monotonic per-process counter (below) -/// still tells two files apart within one process, same as before. -inline unsigned long CurrentProcessId() { -#if defined(_WIN32) - return static_cast(_getpid()); -#else - return static_cast(getpid()); -#endif -} - -/// \brief a counter unique within this process, atomic because nothing here -/// promises the tests that call Write() stay single-threaded. -inline std::uint64_t NextSerial() { - static std::atomic serial{0}; - return serial.fetch_add(1, std::memory_order_relaxed) + 1; -} - -/// \brief the process id encoded in one of Write()'s filenames -/// \return nullopt when `filename` is not one of ours, or carries no pid -/// \note The name is `flm_gemma4_