From 095ca88965f5b374afcd8fca9e705d9be17f5263 Mon Sep 17 00:00:00 2001 From: shuyuzhang Date: Sat, 14 Nov 2026 15:05:07 +0800 Subject: [PATCH 1/4] edge0-strata: run Edge0 145B and Edge0 120B on Metal (Apple Silicon) Restore metal/ from Strata f1933b1 and apply the Metal latent Q4 handoff patch; add the Metal launch templates (edge0-145b-metal.json, edge0-120b-metal.json) with their prepare step, tools/mtp_gguf.py to rebuild mtp-q2_0.gguf from the released mtp/ folder, contract 5, and the macOS build in EDGE0.md. --- edge0-strata/.gitignore | 4 + edge0-strata/EDGE0.md | 89 +- edge0-strata/edge0/CONTRACT | 2 +- .../edge0/models/edge0-120b-metal.json | 70 + .../edge0/models/edge0-145b-metal.json | 70 + edge0-strata/metal/CMakeLists.txt | 187 + .../metal/include/metal_compat/cpuid.h | 6 + .../metal/include/metal_compat/cublas_v2.h | 72 + .../metal/include/metal_compat/cuda_bf16.h | 59 + .../metal_compat/cuda_device_builtins.h | 165 + .../metal/include/metal_compat/cuda_fp16.h | 99 + .../metal/include/metal_compat/cuda_fp8.h | 4 + .../metal/include/metal_compat/cuda_launch.h | 31 + .../metal/include/metal_compat/cuda_runtime.h | 218 ++ .../include/metal_compat/ds4_metal_host.h | 111 + .../metal/include/metal_compat/immintrin.h | 5 + .../metal/include/metal_compat/linux/mman.h | 12 + edge0-strata/metal/include/metal_compat/mma.h | 38 + .../metal/include/metal_compat/vector_types.h | 25 + .../metal/include/strata/gpu/metal.hpp | 44 + .../metal/include/strata/gpu/runtime.hpp | 127 + edge0-strata/metal/kernels/alu_bench.metal | 181 + edge0-strata/metal/kernels/dense_mm.metal | 77 + edge0-strata/metal/kernels/gdn.metal | 185 + edge0-strata/metal/kernels/gdn_prompt.metal | 789 ++++ edge0-strata/metal/kernels/gdn_tprobe.metal | 75 + edge0-strata/metal/kernels/gemv.metal | 264 ++ edge0-strata/metal/kernels/glue.metal | 63 + edge0-strata/metal/kernels/gr.metal | 179 + edge0-strata/metal/kernels/gr2.metal | 341 ++ edge0-strata/metal/kernels/gridbar.metal | 54 + .../metal/kernels/latent_prompt.metal | 54 + .../kernels/latent_prompt_q3_stage.metal | 71 + .../metal/kernels/latent_routed.metal | 181 + edge0-strata/metal/kernels/mm_bench.metal | 85 + edge0-strata/metal/kernels/mma.metal | 216 ++ edge0-strata/metal/kernels/moe.metal | 342 ++ edge0-strata/metal/kernels/moe_mm.metal | 1160 ++++++ edge0-strata/metal/kernels/prompt.metal | 177 + edge0-strata/metal/kernels/qsa.metal | 409 ++ edge0-strata/metal/kernels/qsa2.metal | 290 ++ edge0-strata/metal/kernels/qsa_prompt.metal | 348 ++ edge0-strata/metal/kernels/quant.metal | 701 ++++ edge0-strata/metal/kernels/router.metal | 104 + edge0-strata/metal/kernels/rt_bench.metal | 67 + edge0-strata/metal/kernels/sample.metal | 59 + edge0-strata/metal/kernels/units.metal | 321 ++ edge0-strata/metal/kernels/wrap_probe.metal | 12 + edge0-strata/metal/probe/mkfile.c | 23 + edge0-strata/metal/probe/mtl_doorbell.cpp | 115 + edge0-strata/metal/probe/mtl_probe.cpp | 268 ++ edge0-strata/metal/setup_mac.py | 578 +++ edge0-strata/metal/src/engine/arena_budget.h | 177 + edge0-strata/metal/src/engine/ds4_ctx.cpp | 80 + edge0-strata/metal/src/engine/ds4_finish.cpp | 2754 ++++++++++++++ edge0-strata/metal/src/engine/ds4_host.cpp | 254 ++ .../metal/src/engine/exl3_legacy_stub.cpp | 29 + edge0-strata/metal/src/engine/iq_host.cpp | 98 + .../metal/src/engine/latent_stubs.cpp | 152 + edge0-strata/metal/src/engine/missing.cpp | 80 + edge0-strata/metal/src/engine/mmvq_host.cpp | 66 + .../metal/src/engine/msl/ds4_init.metal | 2800 ++++++++++++++ .../metal/src/engine/msl/ds4_iq.metal | 1184 ++++++ .../metal/src/engine/msl/ds4_mm.metal | 139 + .../metal/src/engine/msl/ds4_quant.metal | 32 + .../metal/src/engine/msl/ggml_grids.h | 21 + edge0-strata/metal/src/engine/msl/mtp.metal | 68 + .../metal/src/engine/msl/mtp_chain.metal | 92 + edge0-strata/metal/src/engine/msl/ple.metal | 103 + .../metal/src/engine/prompt_chain.inc | 571 +++ edge0-strata/metal/src/engine/setup_host.cpp | 160 + edge0-strata/metal/src/engine/stream.cpp | 824 ++++ edge0-strata/metal/src/engine/stream.h | 132 + edge0-strata/metal/src/engine/stubs.cpp | 371 ++ edge0-strata/metal/src/engine/stubs_0145.cpp | 165 + edge0-strata/metal/src/engine/stubs_023.cpp | 15 + edge0-strata/metal/src/engine/stubs_025.cpp | 19 + edge0-strata/metal/src/engine/stubs_144.cpp | 64 + edge0-strata/metal/src/engine/stubs_ds4.cpp | 16 + edge0-strata/metal/src/engine/stubs_gyro.cpp | 25 + edge0-strata/metal/src/engine/window.cpp | 3334 +++++++++++++++++ edge0-strata/metal/src/gguf_lite.cpp | 117 + edge0-strata/metal/src/gguf_lite.h | 37 + edge0-strata/metal/src/gpu/runtime.cpp | 857 +++++ edge0-strata/metal/src/mtl_rt.h | 150 + edge0-strata/metal/src/mtl_rt.mm | 1117 ++++++ edge0-strata/metal/src/ref_quant.cpp | 517 +++ edge0-strata/metal/src/ref_quant.h | 47 + edge0-strata/metal/src/resource_path.h | 56 + edge0-strata/metal/tests/alu_bench.cpp | 39 + .../metal/tests/arena_budget_test.cpp | 135 + edge0-strata/metal/tests/coop_layout.cpp | 32 + edge0-strata/metal/tests/coop_probe.h | 36 + edge0-strata/metal/tests/dense_mm_parity.cpp | 139 + edge0-strata/metal/tests/dsu_job.metal | 45 + edge0-strata/metal/tests/dsu_parity.cpp | 191 + edge0-strata/metal/tests/dsu_ref_vs_ggml.cpp | 29 + .../metal/tests/dsv41/bf16_mv_parity.mm | 133 + .../metal/tests/dsv41/iq2s_t8_parity.mm | 193 + .../metal/tests/dsv41/iq2xs_t8_parity.mm | 190 + .../metal/tests/dsv41/iq2xxs_gemv_parity.mm | 185 + .../metal/tests/dsv41/iq2xxs_t8_parity.mm | 186 + .../metal/tests/dsv41/q8_gemv_parity.mm | 117 + .../metal/tests/dsv41/q8_quant_parity.mm | 105 + edge0-strata/metal/tests/expert_parity.cpp | 669 ++++ edge0-strata/metal/tests/fuse_parity.cpp | 160 + edge0-strata/metal/tests/gdn_chunk_bench.cpp | 374 ++ edge0-strata/metal/tests/gdn_layer_bench.cpp | 194 + edge0-strata/metal/tests/gdn_parity.cpp | 213 ++ .../metal/tests/gdn_prompt_parity.cpp | 175 + edge0-strata/metal/tests/gdn_tprobe.cpp | 36 + edge0-strata/metal/tests/gemv_parity.cpp | 333 ++ edge0-strata/metal/tests/gr2_parity.cpp | 246 ++ edge0-strata/metal/tests/gr_parity.cpp | 208 + edge0-strata/metal/tests/gridbar_bench.cpp | 66 + edge0-strata/metal/tests/gyro_parity.sh | 39 + edge0-strata/metal/tests/mm_bench.cpp | 116 + edge0-strata/metal/tests/mma_parity.cpp | 178 + edge0-strata/metal/tests/moe_mm_parity.cpp | 268 ++ edge0-strata/metal/tests/pair_parity.cpp | 118 + edge0-strata/metal/tests/prompt_chain.cpp | 555 +++ edge0-strata/metal/tests/prompt_parity.cpp | 222 ++ edge0-strata/metal/tests/qsa_parity.cpp | 308 ++ .../metal/tests/qsa_prompt_parity.cpp | 160 + edge0-strata/metal/tests/ref_dump.cpp | 30 + edge0-strata/metal/tests/route_parity.cpp | 90 + edge0-strata/metal/tests/rt_bench.cpp | 238 ++ edge0-strata/metal/tests/runtime_hello.cpp | 162 + edge0-strata/metal/tests/sample_parity.cpp | 61 + edge0-strata/metal/tests/serve_probe.py | 66 + edge0-strata/metal/tests/token_chain.cpp | 661 ++++ edge0-strata/metal/tests/wrap_split_test.cpp | 95 + edge0-strata/metal/tools/chat_tokens.py | 23 + edge0-strata/metal/tools/check_ref.py | 34 + edge0-strata/metal/tools/gen_grids.sh | 29 + edge0-strata/metal/tools/gen_stubs.py | 184 + edge0-strata/metal/tools/gen_tq_lut.py | 50 + edge0-strata/metal/tools/gguf_list.cpp | 25 + edge0-strata/metal/tools/metal_repack.py | 243 ++ edge0-strata/metal/tools/stream_experts.py | 102 + edge0-strata/src/core/verify.cpp | 2 + edge0-strata/src/prefill/prefill.cpp | 2 + edge0-strata/src/program/generate.cpp | 14 + edge0-strata/tools/configure_metal_latent.py | 112 + edge0-strata/tools/mtp_gguf.py | 146 + edge0-strata/tools/verify_metal_latent.py | 89 + 146 files changed, 34795 insertions(+), 6 deletions(-) create mode 100644 edge0-strata/edge0/models/edge0-120b-metal.json create mode 100644 edge0-strata/edge0/models/edge0-145b-metal.json create mode 100644 edge0-strata/metal/CMakeLists.txt create mode 100644 edge0-strata/metal/include/metal_compat/cpuid.h create mode 100644 edge0-strata/metal/include/metal_compat/cublas_v2.h create mode 100644 edge0-strata/metal/include/metal_compat/cuda_bf16.h create mode 100644 edge0-strata/metal/include/metal_compat/cuda_device_builtins.h create mode 100644 edge0-strata/metal/include/metal_compat/cuda_fp16.h create mode 100644 edge0-strata/metal/include/metal_compat/cuda_fp8.h create mode 100644 edge0-strata/metal/include/metal_compat/cuda_launch.h create mode 100644 edge0-strata/metal/include/metal_compat/cuda_runtime.h create mode 100644 edge0-strata/metal/include/metal_compat/ds4_metal_host.h create mode 100644 edge0-strata/metal/include/metal_compat/immintrin.h create mode 100644 edge0-strata/metal/include/metal_compat/linux/mman.h create mode 100644 edge0-strata/metal/include/metal_compat/mma.h create mode 100644 edge0-strata/metal/include/metal_compat/vector_types.h create mode 100644 edge0-strata/metal/include/strata/gpu/metal.hpp create mode 100644 edge0-strata/metal/include/strata/gpu/runtime.hpp create mode 100644 edge0-strata/metal/kernels/alu_bench.metal create mode 100644 edge0-strata/metal/kernels/dense_mm.metal create mode 100644 edge0-strata/metal/kernels/gdn.metal create mode 100644 edge0-strata/metal/kernels/gdn_prompt.metal create mode 100644 edge0-strata/metal/kernels/gdn_tprobe.metal create mode 100644 edge0-strata/metal/kernels/gemv.metal create mode 100644 edge0-strata/metal/kernels/glue.metal create mode 100644 edge0-strata/metal/kernels/gr.metal create mode 100644 edge0-strata/metal/kernels/gr2.metal create mode 100644 edge0-strata/metal/kernels/gridbar.metal create mode 100644 edge0-strata/metal/kernels/latent_prompt.metal create mode 100644 edge0-strata/metal/kernels/latent_prompt_q3_stage.metal create mode 100644 edge0-strata/metal/kernels/latent_routed.metal create mode 100644 edge0-strata/metal/kernels/mm_bench.metal create mode 100644 edge0-strata/metal/kernels/mma.metal create mode 100644 edge0-strata/metal/kernels/moe.metal create mode 100644 edge0-strata/metal/kernels/moe_mm.metal create mode 100644 edge0-strata/metal/kernels/prompt.metal create mode 100644 edge0-strata/metal/kernels/qsa.metal create mode 100644 edge0-strata/metal/kernels/qsa2.metal create mode 100644 edge0-strata/metal/kernels/qsa_prompt.metal create mode 100644 edge0-strata/metal/kernels/quant.metal create mode 100644 edge0-strata/metal/kernels/router.metal create mode 100644 edge0-strata/metal/kernels/rt_bench.metal create mode 100644 edge0-strata/metal/kernels/sample.metal create mode 100644 edge0-strata/metal/kernels/units.metal create mode 100644 edge0-strata/metal/kernels/wrap_probe.metal create mode 100644 edge0-strata/metal/probe/mkfile.c create mode 100644 edge0-strata/metal/probe/mtl_doorbell.cpp create mode 100644 edge0-strata/metal/probe/mtl_probe.cpp create mode 100644 edge0-strata/metal/setup_mac.py create mode 100644 edge0-strata/metal/src/engine/arena_budget.h create mode 100644 edge0-strata/metal/src/engine/ds4_ctx.cpp create mode 100644 edge0-strata/metal/src/engine/ds4_finish.cpp create mode 100644 edge0-strata/metal/src/engine/ds4_host.cpp create mode 100644 edge0-strata/metal/src/engine/exl3_legacy_stub.cpp create mode 100644 edge0-strata/metal/src/engine/iq_host.cpp create mode 100644 edge0-strata/metal/src/engine/latent_stubs.cpp create mode 100644 edge0-strata/metal/src/engine/missing.cpp create mode 100644 edge0-strata/metal/src/engine/mmvq_host.cpp create mode 100644 edge0-strata/metal/src/engine/msl/ds4_init.metal create mode 100644 edge0-strata/metal/src/engine/msl/ds4_iq.metal create mode 100644 edge0-strata/metal/src/engine/msl/ds4_mm.metal create mode 100644 edge0-strata/metal/src/engine/msl/ds4_quant.metal create mode 100644 edge0-strata/metal/src/engine/msl/ggml_grids.h create mode 100644 edge0-strata/metal/src/engine/msl/mtp.metal create mode 100644 edge0-strata/metal/src/engine/msl/mtp_chain.metal create mode 100644 edge0-strata/metal/src/engine/msl/ple.metal create mode 100644 edge0-strata/metal/src/engine/prompt_chain.inc create mode 100644 edge0-strata/metal/src/engine/setup_host.cpp create mode 100644 edge0-strata/metal/src/engine/stream.cpp create mode 100644 edge0-strata/metal/src/engine/stream.h create mode 100644 edge0-strata/metal/src/engine/stubs.cpp create mode 100644 edge0-strata/metal/src/engine/stubs_0145.cpp create mode 100644 edge0-strata/metal/src/engine/stubs_023.cpp create mode 100644 edge0-strata/metal/src/engine/stubs_025.cpp create mode 100644 edge0-strata/metal/src/engine/stubs_144.cpp create mode 100644 edge0-strata/metal/src/engine/stubs_ds4.cpp create mode 100644 edge0-strata/metal/src/engine/stubs_gyro.cpp create mode 100644 edge0-strata/metal/src/engine/window.cpp create mode 100644 edge0-strata/metal/src/gguf_lite.cpp create mode 100644 edge0-strata/metal/src/gguf_lite.h create mode 100644 edge0-strata/metal/src/gpu/runtime.cpp create mode 100644 edge0-strata/metal/src/mtl_rt.h create mode 100644 edge0-strata/metal/src/mtl_rt.mm create mode 100644 edge0-strata/metal/src/ref_quant.cpp create mode 100644 edge0-strata/metal/src/ref_quant.h create mode 100644 edge0-strata/metal/src/resource_path.h create mode 100644 edge0-strata/metal/tests/alu_bench.cpp create mode 100644 edge0-strata/metal/tests/arena_budget_test.cpp create mode 100644 edge0-strata/metal/tests/coop_layout.cpp create mode 100644 edge0-strata/metal/tests/coop_probe.h create mode 100644 edge0-strata/metal/tests/dense_mm_parity.cpp create mode 100644 edge0-strata/metal/tests/dsu_job.metal create mode 100644 edge0-strata/metal/tests/dsu_parity.cpp create mode 100644 edge0-strata/metal/tests/dsu_ref_vs_ggml.cpp create mode 100644 edge0-strata/metal/tests/dsv41/bf16_mv_parity.mm create mode 100644 edge0-strata/metal/tests/dsv41/iq2s_t8_parity.mm create mode 100644 edge0-strata/metal/tests/dsv41/iq2xs_t8_parity.mm create mode 100644 edge0-strata/metal/tests/dsv41/iq2xxs_gemv_parity.mm create mode 100644 edge0-strata/metal/tests/dsv41/iq2xxs_t8_parity.mm create mode 100644 edge0-strata/metal/tests/dsv41/q8_gemv_parity.mm create mode 100644 edge0-strata/metal/tests/dsv41/q8_quant_parity.mm create mode 100644 edge0-strata/metal/tests/expert_parity.cpp create mode 100644 edge0-strata/metal/tests/fuse_parity.cpp create mode 100644 edge0-strata/metal/tests/gdn_chunk_bench.cpp create mode 100644 edge0-strata/metal/tests/gdn_layer_bench.cpp create mode 100644 edge0-strata/metal/tests/gdn_parity.cpp create mode 100644 edge0-strata/metal/tests/gdn_prompt_parity.cpp create mode 100644 edge0-strata/metal/tests/gdn_tprobe.cpp create mode 100644 edge0-strata/metal/tests/gemv_parity.cpp create mode 100644 edge0-strata/metal/tests/gr2_parity.cpp create mode 100644 edge0-strata/metal/tests/gr_parity.cpp create mode 100644 edge0-strata/metal/tests/gridbar_bench.cpp create mode 100644 edge0-strata/metal/tests/gyro_parity.sh create mode 100644 edge0-strata/metal/tests/mm_bench.cpp create mode 100644 edge0-strata/metal/tests/mma_parity.cpp create mode 100644 edge0-strata/metal/tests/moe_mm_parity.cpp create mode 100644 edge0-strata/metal/tests/pair_parity.cpp create mode 100644 edge0-strata/metal/tests/prompt_chain.cpp create mode 100644 edge0-strata/metal/tests/prompt_parity.cpp create mode 100644 edge0-strata/metal/tests/qsa_parity.cpp create mode 100644 edge0-strata/metal/tests/qsa_prompt_parity.cpp create mode 100644 edge0-strata/metal/tests/ref_dump.cpp create mode 100644 edge0-strata/metal/tests/route_parity.cpp create mode 100644 edge0-strata/metal/tests/rt_bench.cpp create mode 100644 edge0-strata/metal/tests/runtime_hello.cpp create mode 100644 edge0-strata/metal/tests/sample_parity.cpp create mode 100644 edge0-strata/metal/tests/serve_probe.py create mode 100644 edge0-strata/metal/tests/token_chain.cpp create mode 100644 edge0-strata/metal/tests/wrap_split_test.cpp create mode 100644 edge0-strata/metal/tools/chat_tokens.py create mode 100644 edge0-strata/metal/tools/check_ref.py create mode 100644 edge0-strata/metal/tools/gen_grids.sh create mode 100644 edge0-strata/metal/tools/gen_stubs.py create mode 100644 edge0-strata/metal/tools/gen_tq_lut.py create mode 100644 edge0-strata/metal/tools/gguf_list.cpp create mode 100644 edge0-strata/metal/tools/metal_repack.py create mode 100644 edge0-strata/metal/tools/stream_experts.py create mode 100644 edge0-strata/tools/configure_metal_latent.py create mode 100644 edge0-strata/tools/mtp_gguf.py create mode 100644 edge0-strata/tools/verify_metal_latent.py diff --git a/edge0-strata/.gitignore b/edge0-strata/.gitignore index 65457fdc..62a37966 100644 --- a/edge0-strata/.gitignore +++ b/edge0-strata/.gitignore @@ -1,3 +1,7 @@ +# Metal build tree and the downloaded Apple metal-cpp headers. +/build-metal/ +/third_party/metal-cpp* + # Generated, not source: deterministic build products of tools/ (see docs/). # tiny-qwen4exp.gguf is ~113 MB, over GitHub's 100 MB file limit -- regenerate with # python tools/make_tiny_model.py diff --git a/edge0-strata/EDGE0.md b/edge0-strata/EDGE0.md index 4ec3743e..f0c1574c 100644 --- a/edge0-strata/EDGE0.md +++ b/edge0-strata/EDGE0.md @@ -18,6 +18,8 @@ start it, wait until it is ready and talk to it on a loopback port. See - Imported from the Strata fork in which latent-MoE support for these models was developed, at commit `f1933b1e4993fecc10d2e9360d85569619a2c42d`, together with source files that were not yet committed there; their paths and SHA256 checksums are recorded in [`edge0/IMPORTED_UNCOMMITTED.txt`](edge0/IMPORTED_UNCOMMITTED.txt). +- `metal/` is that same commit's `metal/` directory, restored byte for byte (`git archive f1933b1 metal`). + It is not part of the uncommitted-file list above. - The fork's merge base with upstream Strata (https://github.com/Niko1221/Strata) is `c62d5e3a15b40ea51c96aae9e8d36928965f53ab` (2026-10-31), a public upstream commit. @@ -56,7 +58,7 @@ In this directory: - Removed because neither the Linux CUDA build nor the server uses them (each was checked against the CMake files and the Python imports before removal): - experiment files in the root directory (`q4du_utest*`, `run-latent-du.sh`); - - the other engines and their setup: `sycl/` (Intel Arc), `metal/` (Apple Silicon); + - the other engines and their setup: `sycl/` (Intel Arc); - the desktop app (`apps/`) and the release packaging scripts (`packaging/`; this directory is built with CMake as described below); - the container image (`docker/`) and the NVIDIA DGX Spark engine build (`spark/`); @@ -193,10 +195,69 @@ interpreter. To run `build\strata.exe` by hand, put the CUDA DLLs on `PATH` firs (`set "PATH=%CUDA_PATH%\bin;%PATH%"`). The engine writes temporary files under `%TMP%`, which must be writable for the user who runs `edge0-infer`. +## Build (macOS, Apple Silicon) + +Requirements: an Apple Silicon Mac, the Xcode Command Line Tools, CMake 3.24 or newer, Ninja, Python +3.10 or newer, and [uv](https://docs.astral.sh/uv/). Run the commands below in this directory +(`edge0-strata/`). `build-metal/`, `.venv/` and the downloaded headers are ignored by Git. + +1. Unpack the same pinned ggml source as on Linux into `third_party/llama.cpp`: + + ``` + mkdir -p third_party/llama.cpp + curl -L https://github.com/ggml-org/llama.cpp/archive/3cf03257f219afbe7334045ff7c6a06ac68c627d.tar.gz \ + | tar -xz --strip-components=1 -C third_party/llama.cpp + ``` + +2. Download Apple's metal-cpp 26 headers and check the SHA-256 recorded in `metal/setup_mac.py` + (`4df3c078b9aadcb516212e9cb03004cbc5ce9a3e9c068fa3144d021db585a3a4`). The archive contains a + `metal-cpp/` directory; unpack it to `third_party/metal-cpp`: + + ``` + mkdir -p third_party + curl -L -o third_party/metal-cpp_26.zip https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip + echo "4df3c078b9aadcb516212e9cb03004cbc5ce9a3e9c068fa3144d021db585a3a4 third_party/metal-cpp_26.zip" | shasum -a 256 -c + rm -rf third_party/metal-cpp + unzip -q third_party/metal-cpp_26.zip -d third_party + ``` + + `third_party/metal-cpp/Metal/Metal.hpp` is present after a successful unpack. + +3. Create `.venv` and install the server packages and the build tools (CMake and Ninja are in + `requirements-build.txt`): + + ``` + uv venv .venv + uv pip install --python .venv/bin/python -r requirements.txt -r requirements-build.txt + ``` + +4. Build the engine. Put the virtual environment's `bin` on `PATH` so CMake finds that Ninja: + + ``` + export PATH="$PWD/.venv/bin:$PATH" + cmake -S metal -B build-metal -G Ninja -DCMAKE_BUILD_TYPE=Release \ + -DMETAL_CPP_DIR="$PWD/third_party/metal-cpp" \ + -DSTRATA_GGML_DIR="$PWD/third_party/llama.cpp" + cmake --build build-metal --target strata + cmake --build build-metal --target moe_mm_parity + ``` + + The engine is `build-metal/strata`. The operator check is `build-metal/moe_mm_parity`; both of + these print `PASS`: + + ``` + ./build-metal/moe_mm_parity --gu Q3_K --dn IQ4_NL --latent --v2 --stage-gu --tile 32,64,32,2 --T 64 + ./build-metal/moe_mm_parity --gu Q3_K --dn IQ4_NL --latent --v2 --stage-gu --tile 32,64,32,2 --T 64 --latent1792 + ``` + +`edge0-infer` on macOS starts `build-metal/strata` through `.venv/bin/python` and `serve/server.py`. +The device is `metal`, which is also the default, and the only accepted value. + ## Launch contract with edge0-infer -- `edge0/CONTRACT` holds one integer, the contract version (currently 4). `edge0-infer` refuses to start an - engine whose version it does not support. +- `edge0/CONTRACT` holds one integer, the contract version (currently 5). `edge0-infer` refuses to start an + engine whose version it does not support. Version 5 is version 4 plus the Metal template files and the + optional `prepare` field described below. - `edge0/models/.json` is a Strata server configuration template for each delegated model. `edge0-infer` replaces the placeholders and passes the result to `serve/server.py --engine strata --config --host 127.0.0.1 --port 0 --ready-file `: @@ -215,8 +276,26 @@ writable for the user who runs `edge0-infer`. (Edge0 120B, 1280 latent dimensions). `edge0-infer` chooses the D/U default (Q8 for Edge0 145B, Q4 for Edge0 120B) and the expert format, which it reads from `model_name` in the folder's `model.json`. - `edge0-infer` also sets `gpu` from `--device cuda:`, `model_name` from `--alias`, and the - `--max-context` argument from `--context`. + On macOS the device is `metal`. `edge0-infer` then loads `edge0-145b-metal.json` or + `edge0-120b-metal.json` instead of the CUDA template. Those files use the same placeholders. Their + `lib_dirs` list is empty, so `${CUDA_LIB_DIR}` is not used. They set `--max-context` to `32768` and + `--spec` to `4`, and they keep `STRATA_DECODE_TIMING=1`. They do not set `STRATA_METAL_MTP_GGUF`: + the engine reads `/mtp-q2_0.gguf`. `STRATA_METAL_DRAFT_VOCAB` points at + `${MODEL_DIR}/mtp/draft_vocab.bin`. + + A Metal template may contain `prepare`, an array of files to create before the server starts. Each + entry has `output`, `script` (a path relative to this directory) and `args`, and those strings use + the placeholders above. `edge0-infer` runs a missing `output` with this directory's Python, or with + the interpreter given by `--strata-python` when that option is set, and with this directory as the + working directory. It then removes `prepare` from the configuration it passes to + `server.py`. The shipped entries run `tools/mtp_gguf.py` to write `mtp-q2_0.gguf` into the model + directory, so that directory must be writable. An `output` that already exists is left as it is. + + Metal accepts only Q3_K expert folders (`model_name` ending in `-q3k`), `du=q4` and + `ple-io=direct`. An NVFP4 folder, `du=q8` or `ple-io=ram` is rejected before any process starts. + + `edge0-infer` also sets `gpu` from `--device` (`cuda:` on Linux and Windows, `metal` on macOS), + `model_name` from `--alias`, and the `--max-context` argument from `--context`. A change to the placeholders, the command line or the readiness protocol increments the contract version. diff --git a/edge0-strata/edge0/CONTRACT b/edge0-strata/edge0/CONTRACT index b8626c4c..7ed6ff82 100644 --- a/edge0-strata/edge0/CONTRACT +++ b/edge0-strata/edge0/CONTRACT @@ -1 +1 @@ -4 +5 diff --git a/edge0-strata/edge0/models/edge0-120b-metal.json b/edge0-strata/edge0/models/edge0-120b-metal.json new file mode 100644 index 00000000..e28b19c4 --- /dev/null +++ b/edge0-strata/edge0/models/edge0-120b-metal.json @@ -0,0 +1,70 @@ +{ + "exe": "${STRATA_DIR}/build-metal/strata", + "cwd": "${STRATA_DIR}", + "backend": "metal", + "args": [ + "--pack", + "${MODEL_DIR}", + "--native", + "${MODEL_DIR}/model-skeleton.gguf", + "--native-dense-gguf", + "${MODEL_DIR}/model-skeleton.gguf", + "--ple-gguf", + "${MODEL_DIR}/ple-fp8.gguf", + "--ple-io", + "${PLE_IO}", + "--expert-profile", + "${MODEL_DIR}/profile.bin", + "--expert-cache", + "64", + "--mmap-experts", + "--prefill", + "auto", + "--max-context", + "32768", + "--spec", + "4", + "--suffix-draft", + "3", + "--mtp", + "${MODEL_DIR}/mtp", + "--spec-min-p", + "0.5" + ], + "env": { + "STRATA_METAL_LATENT": "1", + "STRATA_LATENT_DENSE": "${DU_FILE}", + "STRATA_LATENT_DENSE_FMT": "${DU_FMT}", + "STRATA_LATENT_A": "${MODEL_DIR}/experts.bin", + "STRATA_LATENT_LABELS": "${MODEL_DIR}/labels.bin", + "STRATA_HC_Q8": "0", + "STRATA_ROUTE_TAIL_SKIP": "0", + "STRATA_PREFILL_CPU_SHARE": "0", + "STRATA_METAL_LATENT_GROUP_U": "1", + "STRATA_METAL_LATENT_PROMPT": "1", + "STRATA_METAL_PROMPT_CHUNK": "8192", + "STRATA_METAL_LATENT_DECODE_TILES": "1", + "STRATA_METAL_MATH": "fast", + "STRATA_METAL_LATENT_PROMPT_QSA_MM": "1", + "STRATA_METAL_GDN_SCAN": "gdn_scanr_4_8", + "STRATA_METAL_DRAFT_VOCAB": "${MODEL_DIR}/mtp/draft_vocab.bin", + "STRATA_DECODE_TIMING": "1", + "STRATA_METAL_LATENT_VERIFY_FAST": "1", + "STRATA_METAL_LATENT_PROMPT_TUNE": "1" + }, + "tokenizer": "${MODEL_DIR}/tokenizer", + "model_name": "edge0-120b", + "lib_dirs": [], + "prepare": [ + { + "output": "${MODEL_DIR}/mtp-q2_0.gguf", + "script": "tools/mtp_gguf.py", + "args": [ + "--rt", + "${MODEL_DIR}/mtp", + "--out", + "${MODEL_DIR}/mtp-q2_0.gguf" + ] + } + ] +} diff --git a/edge0-strata/edge0/models/edge0-145b-metal.json b/edge0-strata/edge0/models/edge0-145b-metal.json new file mode 100644 index 00000000..2cabc7b4 --- /dev/null +++ b/edge0-strata/edge0/models/edge0-145b-metal.json @@ -0,0 +1,70 @@ +{ + "exe": "${STRATA_DIR}/build-metal/strata", + "cwd": "${STRATA_DIR}", + "backend": "metal", + "args": [ + "--pack", + "${MODEL_DIR}", + "--native", + "${MODEL_DIR}/model-skeleton.gguf", + "--native-dense-gguf", + "${MODEL_DIR}/model-skeleton.gguf", + "--ple-gguf", + "${MODEL_DIR}/ple-fp8.gguf", + "--ple-io", + "${PLE_IO}", + "--expert-profile", + "${MODEL_DIR}/profile.bin", + "--expert-cache", + "64", + "--mmap-experts", + "--prefill", + "auto", + "--max-context", + "32768", + "--spec", + "4", + "--suffix-draft", + "3", + "--mtp", + "${MODEL_DIR}/mtp", + "--spec-min-p", + "0.5" + ], + "env": { + "STRATA_METAL_LATENT": "1", + "STRATA_LATENT_DENSE": "${DU_FILE}", + "STRATA_LATENT_DENSE_FMT": "${DU_FMT}", + "STRATA_LATENT_A": "${MODEL_DIR}/experts.bin", + "STRATA_LATENT_LABELS": "${MODEL_DIR}/labels.bin", + "STRATA_HC_Q8": "0", + "STRATA_ROUTE_TAIL_SKIP": "0", + "STRATA_PREFILL_CPU_SHARE": "0", + "STRATA_METAL_LATENT_GROUP_U": "1", + "STRATA_METAL_LATENT_PROMPT": "1", + "STRATA_METAL_PROMPT_CHUNK": "8192", + "STRATA_METAL_LATENT_DECODE_TILES": "1", + "STRATA_METAL_MATH": "fast", + "STRATA_METAL_LATENT_PROMPT_QSA_MM": "1", + "STRATA_METAL_GDN_SCAN": "gdn_scanr_4_8", + "STRATA_METAL_DRAFT_VOCAB": "${MODEL_DIR}/mtp/draft_vocab.bin", + "STRATA_DECODE_TIMING": "1", + "STRATA_METAL_LATENT_VERIFY_FAST": "1", + "STRATA_METAL_LATENT_PROMPT_TUNE": "1" + }, + "tokenizer": "${MODEL_DIR}/tokenizer", + "model_name": "edge0-145b", + "lib_dirs": [], + "prepare": [ + { + "output": "${MODEL_DIR}/mtp-q2_0.gguf", + "script": "tools/mtp_gguf.py", + "args": [ + "--rt", + "${MODEL_DIR}/mtp", + "--out", + "${MODEL_DIR}/mtp-q2_0.gguf" + ] + } + ] +} diff --git a/edge0-strata/metal/CMakeLists.txt b/edge0-strata/metal/CMakeLists.txt new file mode 100644 index 00000000..4ddf36d6 --- /dev/null +++ b/edge0-strata/metal/CMakeLists.txt @@ -0,0 +1,187 @@ +# metal/ - Strata on Apple Silicon (Metal), MAC-PLAN.md. Its own CMake project like sycl/, macOS only, never part of +# the CUDA / HIP / SYCL builds. +# +# Two parts: +# * the kernel tree (mac-kern): the runtime face src/mtl_rt.*, the MSL kernels kernels/*.metal (compiled at run time), +# CPU references src/ref_*.cpp and the parity / bench programs tests/*.cpp; +# * the engine (mac-design): the shared host sources from ../src, UNCHANGED - include/metal_compat goes first on the +# include path and gives them the CUDA runtime names on Metal, as include/strata/hip_compat does on HIP - over +# the Metal device runtime src/gpu/runtime.cpp (metal-cpp) and the kernel library src/engine/*. +# +# cmake -S metal -B build-metal -G Ninja -DCMAKE_BUILD_TYPE=Release -DMETAL_CPP_DIR=... [-DSTRATA_GGML_DIR=...] +# ninja -C build-metal +cmake_minimum_required(VERSION 3.24) +project(strata_metal LANGUAGES C CXX OBJCXX) + +if(NOT APPLE) + message(FATAL_ERROR "metal/ builds on macOS only") +endif() + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_OBJCXX_STANDARD 20) +if(NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE Release) +endif() +set(CMAKE_CXX_FLAGS_RELEASE "-O2") +set(CMAKE_OBJCXX_FLAGS_RELEASE "-O2") + +get_filename_component(STRATA_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/.." ABSOLUTE) +set(GGML_COMMON "${STRATA_ROOT}/third_party/ggml/ggml-common.h") + +# ---- the kernel tree (mac-kern) ------------------------------------------------------------------------------------ +add_library(mtl_rt STATIC src/mtl_rt.mm src/ref_quant.cpp src/gguf_lite.cpp) +target_include_directories(mtl_rt PUBLIC src ${STRATA_ROOT}/third_party/ggml ${STRATA_ROOT}/include) +target_compile_options(mtl_rt PRIVATE $<$:-fobjc-arc>) +target_compile_definitions(mtl_rt PRIVATE + STRATA_GGML_COMMON_DEFAULT="${GGML_COMMON}" + STRATA_METAL_KERNELS_DEFAULT="${CMAKE_CURRENT_SOURCE_DIR}/kernels") +target_link_libraries(mtl_rt PUBLIC "-framework Metal" "-framework Foundation") + +file(GLOB STRATA_METAL_TESTS CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/tests/*.cpp) +list(REMOVE_ITEM STRATA_METAL_TESTS ${CMAKE_CURRENT_SOURCE_DIR}/tests/runtime_hello.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/tests/dsu_ref_vs_ggml.cpp) # (needs ggml: below) +foreach(src ${STRATA_METAL_TESTS}) + get_filename_component(t ${src} NAME_WE) + add_executable(${t} ${src}) + target_link_libraries(${t} PRIVATE mtl_rt) +endforeach() + +# ---- the device runtime and the engine (mac-design) --------------------------------------------------------------- +set(METAL_CPP_DIR "" CACHE PATH "metal-cpp headers (https://developer.apple.com/metal/cpp/), e.g. metal-cpp_26") +if(NOT EXISTS "${METAL_CPP_DIR}/Metal/Metal.hpp") + message(STATUS "strata metal: METAL_CPP_DIR is not set: the device runtime and the engine are not built") + return() +endif() + +add_library(strata_gpu_metal STATIC src/gpu/runtime.cpp) +target_include_directories(strata_gpu_metal BEFORE PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include/metal_compat + ${CMAKE_CURRENT_SOURCE_DIR}/include ${STRATA_ROOT}/include) +target_include_directories(strata_gpu_metal SYSTEM PUBLIC ${METAL_CPP_DIR}) +target_compile_definitions(strata_gpu_metal PUBLIC STRATA_GPU_METAL=1) +target_link_libraries(strata_gpu_metal PUBLIC "-framework Metal" "-framework Foundation" "-framework QuartzCore") + +add_library(strata_core_metal STATIC ${STRATA_ROOT}/src/core/graph.cpp) +target_link_libraries(strata_core_metal PUBLIC strata_gpu_metal) +add_executable(runtime_hello tests/runtime_hello.cpp) +target_link_libraries(runtime_hello PRIVATE strata_core_metal) +enable_testing() +add_test(NAME runtime_hello COMMAND runtime_hello) + +# The engine: ggml (the native expert formats' type traits and ggml-cpu) from a llama.cpp checkout at the commit the +# PC build pins (CMakeLists.txt: 3cf03257). +option(STRATA_METAL_ENGINE "Build the engine (strata) for Metal" ON) +set(STRATA_GGML_DIR "" CACHE PATH "llama.cpp checkout at the pinned commit") +if(NOT STRATA_METAL_ENGINE) + return() +endif() +if(NOT EXISTS "${STRATA_GGML_DIR}/ggml/CMakeLists.txt") + message(FATAL_ERROR "STRATA_GGML_DIR='${STRATA_GGML_DIR}' is not a llama.cpp checkout") +endif() +set(BUILD_SHARED_LIBS OFF CACHE BOOL "" FORCE) +set(GGML_STATIC ON CACHE BOOL "" FORCE) +set(GGML_METAL OFF CACHE BOOL "" FORCE) +set(GGML_BLAS OFF CACHE BOOL "" FORCE) +set(GGML_OPENMP OFF CACHE BOOL "" FORCE) +set(GGML_BACKEND_DL OFF CACHE BOOL "" FORCE) +set(GGML_BUILD_TESTS OFF CACHE BOOL "" FORCE) +set(GGML_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE) +add_subdirectory(${STRATA_GGML_DIR}/ggml ${CMAKE_BINARY_DIR}/ggml EXCLUDE_FROM_ALL) +# ref_quant's IQ1_S / IQ1_M / Q6_K dequantizers against ggml's own (the Unsloth sets' kernels are checked against ref_quant: dsu_parity) +add_executable(dsu_ref_vs_ggml tests/dsu_ref_vs_ggml.cpp) +target_include_directories(dsu_ref_vs_ggml PRIVATE ${STRATA_GGML_DIR}/ggml/include) +target_link_libraries(dsu_ref_vs_ggml PRIVATE mtl_rt ggml-base ggml-cpu) + +set(S ${STRATA_ROOT}/src) +set(STRATA_METAL_HOST_SOURCES + ${S}/core/device.cu ${S}/core/pinned.cu ${S}/platform/memory.cpp ${S}/core/graph.cpp ${S}/core/weights.cpp ${S}/core/layout.cpp + ${S}/core/layer.cpp ${S}/core/session.cpp ${S}/core/expert_source.cpp ${S}/core/remote_experts.cpp ${S}/core/expert_cache.cpp + ${S}/core/native_head.cpp ${S}/core/native_dense.cpp ${S}/core/latent_dense.cpp ${S}/core/verify.cpp ${S}/core/mtp.cpp + ${S}/core/conversation_snapshot.cpp ${S}/core/conversation_state.cpp ${S}/core/conversation_memory.cpp + ${S}/core/peer_experts.cpp ${S}/core/arch_defaults.cpp ${S}/core/vmm.cpp ${S}/core/conversation_file.cpp + ${S}/kernels/ngram.cpp ${S}/ngram/ple_reader.cpp ${S}/platform/direct_file.cpp + ${S}/prefill/prefill.cpp + ${S}/spec/suffix_drafter.cpp ${S}/spec/controller.cpp ${S}/spec/draft_policy.cpp ${S}/spec/draft_source.cpp + ${S}/kernels/cpu/pool.cpp ${S}/kernels/cpu/expert_layout.cpp ${S}/kernels/cpu/native_expert.cpp + ${S}/kernels/cpu/hadamard.cpp ${S}/kernels/cpu/tq_cpu.cpp) # Gyro: the Hadamard rotation, the scalar trellis rows +set_source_files_properties(${S}/core/device.cu ${S}/core/pinned.cu ${DS4}/ds4_vision.cu ${DS4}/vendor/exl3.cu ${DS4}/vendor/moe_win.cu ${DS4}/vendor/sampler.cu + PROPERTIES LANGUAGE CXX) + +# The kernel library: Metal implementations (src/engine/*.cpp) and, for every entry point not ported yet, a stub +# that stops the engine with the entry point's name (src/engine/stubs.cpp, generated by tools/gen_stubs.py). +file(GLOB STRATA_METAL_ENGINE_SOURCES CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/src/engine/*.cpp) + +add_library(strata_metal_host STATIC ${STRATA_METAL_HOST_SOURCES}) +target_include_directories(strata_metal_host PUBLIC ${STRATA_ROOT}/include ${S}/core ${STRATA_ROOT}/third_party/ggml + ${STRATA_GGML_DIR}/ggml/include ${STRATA_GGML_DIR}/ggml/src ${STRATA_GGML_DIR}/ggml/src/ggml-cpu) +target_compile_definitions(strata_metal_host PUBLIC STRATA_NATIVE_EXPERTS=1 STRATA_VERSION="0.1.55-metal") +target_compile_options(strata_metal_host PRIVATE -Wno-unused-parameter -Wno-deprecated-declarations) +target_link_libraries(strata_metal_host PUBLIC strata_gpu_metal ggml-cpu ggml-base) + +add_library(strata_metal_kernels STATIC ${STRATA_METAL_ENGINE_SOURCES}) +target_include_directories(strata_metal_kernels PRIVATE src ${STRATA_ROOT}/third_party/ggml) +target_compile_definitions(strata_metal_kernels PRIVATE + STRATA_GGML_COMMON_DEFAULT="${GGML_COMMON}" STRATA_METAL_KERNELS_DEFAULT="${CMAKE_CURRENT_SOURCE_DIR}/kernels" + STRATA_METAL_ENGINE_KERNELS_DEFAULT="${CMAKE_CURRENT_SOURCE_DIR}/src/engine/msl" + STRATA_METAL_DATA_DEFAULT="${STRATA_ROOT}/data") +target_link_libraries(strata_metal_kernels PUBLIC strata_metal_host mtl_rt) + +add_executable(strata ${S}/program/generate.cpp) +target_link_libraries(strata PRIVATE strata_metal_kernels strata_metal_host strata_metal_kernels) + +# --- the DeepSeek engine on this Metal runtime (M5MAX-DEEPSEEK-PLAN.md, steps 3+). The .cu sources are compiled as C++ +# through include/metal_compat, which gives them CUDA's names over the Metal runtime; the kernel entry points they call are +# the Metal implementations in src/engine/* or the stubs from tools/gen_stubs.py that stop the engine with the name. +set(DS4 ${STRATA_ROOT}/src/dsv41) +# The DeepSeek engine's sources. Every .cu compiles as C++ through include/metal_compat; the EXL3 legacy path is left out +# (the 2.25bpw text model never calls it) and the vision tower is included because ds4_model.cu references it. +set(STRATA_DS4_SOURCES + ${DS4}/ds4_main.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/engine/ds4_finish.cpp ${DS4}/ds4_model.cu ${DS4}/ds4_gguf.cu ${DS4}/ds4_iq.cu ${DS4}/ds4_kernels.cu ${DS4}/ds4_tq.cu + ${DS4}/ds4_vision.cu + ${DS4}/vendor/exl3.cu ${DS4}/vendor/exl3_cpu.cpp ${DS4}/vendor/moe_win.cu ${DS4}/vendor/sampler.cu + ${DS4}/ds4_gg_cpu.cpp ${DS4}/ds4_gguf_cpu.cpp ${DS4}/ds4_iq_cpu.cpp ${DS4}/ds4_tq_cpu.cpp ${DS4}/ds4_packinfo.cpp) +set_source_files_properties( + ${DS4}/ds4_model.cu ${DS4}/ds4_gguf.cu ${DS4}/ds4_iq.cu ${DS4}/ds4_kernels.cu ${DS4}/ds4_tq.cu ${DS4}/ds4_vision.cu + ${DS4}/vendor/exl3.cu ${DS4}/vendor/moe_win.cu ${DS4}/vendor/sampler.cu + PROPERTIES LANGUAGE CXX) +list(APPEND STRATA_DS4_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/src/engine/exl3_legacy_stub.cpp) # v1.1.8 EXL3: absent on Metal +# The CPU lane for the experts the GPU does not hold (ggml-cpu's NEON row dots): without it every non-resident expert contributes +# zero and the residual stream stays the embedding (measured: layers 0-3 changed nothing, then random tokens). +set_source_files_properties(${DS4}/ds4_gg_cpu.cpp PROPERTIES COMPILE_DEFINITIONS "STRATA_DSV41_TQ_CPU=1") +# The DeepSeek engine's queue, residency and barrier rules (deferred command buffers, the second queue, the residency set shared with +# strata::gpu, ...) live in mtl_rt.mm and gpu/runtime.cpp behind STRATA_DS4_RUNTIME. strata_ds4 compiles its own copy of these two +# files with it; the Qwen engine (strata) and every other target keep the code of 0.2.7 (the macro is off there, so the files read +# exactly as they did). The copies are linked as objects, so they win over the same symbols in the libraries. +add_library(ds4_rt_mm OBJECT src/mtl_rt.mm) +target_include_directories(ds4_rt_mm PRIVATE src ${STRATA_ROOT}/third_party/ggml ${STRATA_ROOT}/include) +target_compile_options(ds4_rt_mm PRIVATE $<$:-fobjc-arc>) +target_compile_definitions(ds4_rt_mm PRIVATE STRATA_DS4_RUNTIME=1 STRATA_GGML_COMMON_DEFAULT="${GGML_COMMON}" + STRATA_METAL_KERNELS_DEFAULT="${CMAKE_CURRENT_SOURCE_DIR}/kernels") +add_library(ds4_rt_gpu OBJECT src/gpu/runtime.cpp) +target_include_directories(ds4_rt_gpu BEFORE PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include/metal_compat + ${CMAKE_CURRENT_SOURCE_DIR}/include ${STRATA_ROOT}/include) +target_include_directories(ds4_rt_gpu SYSTEM PRIVATE ${METAL_CPP_DIR}) +target_compile_definitions(ds4_rt_gpu PRIVATE STRATA_DS4_RUNTIME=1 STRATA_GPU_METAL=1) +add_executable(strata_ds4 ${STRATA_DS4_SOURCES} $ $) +target_include_directories(strata_ds4 PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include/metal_compat ${STRATA_ROOT}/include ${DS4} ${CMAKE_CURRENT_SOURCE_DIR}/src + ${STRATA_ROOT}/third_party/ggml ${STRATA_GGML_DIR}/ggml/include ${STRATA_GGML_DIR}/ggml/src ${STRATA_GGML_DIR}/ggml/src/ggml-cpu) +target_compile_definitions(strata_ds4 PRIVATE STRATA_METAL=1) # the real host CPU lane (ds4_gg_cpu.cpp) +target_compile_options(strata_ds4 PRIVATE -Wno-unused-parameter -Wno-deprecated-declarations -Wno-implicit-function-declaration) +target_link_libraries(strata_ds4 PRIVATE strata_metal_kernels strata_metal_host) +# The commit the binary was built from, readable with --version (the root CMakeLists does the same for the CUDA and HIP builds): +# -DSTRATA_BUILD_ID=..., else BUILD_ID beside the sources (a git archive has no .git), else git, else "unknown". +if(NOT STRATA_BUILD_ID AND EXISTS ${STRATA_ROOT}/BUILD_ID) + file(STRINGS ${STRATA_ROOT}/BUILD_ID STRATA_BUILD_ID LIMIT_COUNT 1) +endif() +if(NOT STRATA_BUILD_ID) + execute_process(COMMAND git -C ${STRATA_ROOT} rev-parse --short=9 HEAD OUTPUT_VARIABLE STRATA_BUILD_ID OUTPUT_STRIP_TRAILING_WHITESPACE ERROR_QUIET) +endif() +if(NOT STRATA_BUILD_ID) + set(STRATA_BUILD_ID unknown) +endif() +set_source_files_properties(${DS4}/ds4_main.cpp PROPERTIES COMPILE_DEFINITIONS "DS4_BUILD_ID=\"${STRATA_BUILD_ID}\"") +# -DSTRATA_DS4_VERSION=0.1.55-metal: the version this DeepSeek engine reports, when its package is released apart from the Qwen engine's +if(STRATA_DS4_VERSION) + set_source_files_properties(${DS4}/ds4_main.cpp PROPERTIES COMPILE_OPTIONS "-USTRATA_VERSION;-DSTRATA_VERSION=\"${STRATA_DS4_VERSION}\"") +endif() diff --git a/edge0-strata/metal/include/metal_compat/cpuid.h b/edge0-strata/metal/include/metal_compat/cpuid.h new file mode 100644 index 00000000..8fd3eb52 --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/cpuid.h @@ -0,0 +1,6 @@ +// metal/include/metal_compat/cpuid.h - GCC's on arm64: every leaf reads as zero, so the x86 feature +// probes (src/kernels/cpu/expert_layout.cpp) report no AVX / AVX-512 / VNNI and the engine never picks an x86 kernel. +#pragma once +#define __cpuid(leaf, a, b, c, d) ((void) (leaf), (a) = 0, (b) = 0, (c) = 0, (d) = 0) +#define __cpuid_count(leaf, sub, a, b, c, d) ((void) (leaf), (void) (sub), (a) = 0, (b) = 0, (c) = 0, (d) = 0) +static inline unsigned __get_cpuid_max(unsigned, unsigned* sig) { if (sig) *sig = 0; return 0; } diff --git a/edge0-strata/metal/include/metal_compat/cublas_v2.h b/edge0-strata/metal/include/metal_compat/cublas_v2.h new file mode 100644 index 00000000..c5ebf75e --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/cublas_v2.h @@ -0,0 +1,72 @@ +// metal/include/metal_compat/cublas_v2.h - the cuBLAS surface the dsv41 sources use, over the Metal device runtime. +// +// The engine makes 9 GEMM calls (7 cublasGemmEx, 2 cublasSgemm) for the dense projections, the LM head and the vision +// tower's matmuls. On Metal the dense path is kernels/dense_mm.metal; this header gives the calls the CUDA names so the +// shared sources compile unchanged, and routes them to strata::gpu::dense_gemm(), which the metal host runtime provides. +// +// Compile-audit finding (M5 Max, 2026-10-23): the dsv41 core needs exactly cublasCreate/Destroy/SetStream/SetWorkspace, +// cublasSgemm and cublasGemmEx, plus cublasStatus_t/cublasHandle_t/cublasOperation_t/cublasComputeType_t/cudaDataType_t. +#pragma once + +#include +#include + +#include "cuda_runtime.h" +#include "ds4_metal_host.h" // the kernels that have a Metal dispatch + +typedef void* cublasHandle_t; + +enum cublasStatus_t { CUBLAS_STATUS_SUCCESS = 0, CUBLAS_STATUS_NOT_INITIALIZED = 1, CUBLAS_STATUS_ALLOC_FAILED = 3, + CUBLAS_STATUS_INVALID_VALUE = 7, CUBLAS_STATUS_ARCH_MISMATCH = 8, CUBLAS_STATUS_NOT_SUPPORTED = 14 }; +enum cublasOperation_t { CUBLAS_OP_N = 0, CUBLAS_OP_T = 1, CUBLAS_OP_C = 2 }; +enum cublasComputeType_t { CUBLAS_COMPUTE_16F = 64, CUBLAS_COMPUTE_32F = 68, CUBLAS_COMPUTE_32F_FAST_16F = 74 }; +enum cublasGemmAlgo_t { CUBLAS_GEMM_DEFAULT = -1, CUBLAS_GEMM_DEFAULT_TENSOR_OP = 99 }; +enum cudaDataType_t { CUDA_R_32F = 0, CUDA_R_16F = 2, CUDA_R_16BF = 14, CUDA_R_8F_E4M3 = 28 }; +using cudaDataType = cudaDataType_t; // the engine names it both ways + +static inline cublasStatus_t cublasCreate(cublasHandle_t* h) { + *h = (void*) 1; // the Metal runtime is a singleton: the handle carries no state + return CUBLAS_STATUS_SUCCESS; +} +static inline cublasStatus_t cublasDestroy(cublasHandle_t) { return CUBLAS_STATUS_SUCCESS; } +static inline cublasStatus_t cublasSetStream(cublasHandle_t, cudaStream_t) { return CUBLAS_STATUS_SUCCESS; } +static inline cublasStatus_t cublasSetWorkspace(cublasHandle_t, void*, size_t) { return CUBLAS_STATUS_SUCCESS; } + +// C = alpha * op(A) * op(B) + beta * C, all row-major floats. +// NOT WIRED YET: the metal host runtime has no dense GEMM entry point. This aborts with a message instead of returning a +// wrong result, so the shared sources compile and a run fails where the work is, not somewhere confusing. The real +// implementation drives kernels/dense_mm.metal (the kernel exists and passes its parity test). +[[noreturn]] static inline void strata_metal_dense_gemm_missing(const char* what) { + std::fprintf(stderr, "strata metal: %s is not wired to kernels/dense_mm.metal yet\n", what); + std::abort(); +} + +static inline cublasStatus_t cublasSgemm(cublasHandle_t, cublasOperation_t opA, cublasOperation_t opB, int m, int n, int k, + const float* alpha, const float* A, int lda, const float* B, int ldb, + const float* beta, float* C, int ldc) { + // the mHC mixes (ds4_model.cu:2737): C = A^T B, F32, alpha 1, beta 0 + if (opA == CUBLAS_OP_T && opB == CUBLAS_OP_N && alpha && beta && *alpha == 1.0f && *beta == 0.0f && + strata::dsv41::metal_host::sgemm_tn(A, B, C, m, n, k, lda, ldb, ldc)) { + return CUBLAS_STATUS_SUCCESS; + } + strata_metal_dense_gemm_missing("cublasSgemm"); +} + +// the mixed-precision form: any argument list, because the engine passes cuBLAS's full 19-argument spelling; the body +// names kernels/dense_mm.metal, which is where the real implementation goes. +// the chunk path's call (F16 weights [m][k] x F16 rows [n][k] -> F32 rows [n][m]): the MPP tile GEMM +static inline cublasStatus_t cublasGemmEx(cublasHandle_t, cublasOperation_t ta, cublasOperation_t tb, int m, int n, int k, + const void* alpha, const void* A, cudaDataType_t at, int lda, const void* B, + cudaDataType_t bt, int ldb, const void* beta, void* C, cudaDataType_t ct, int ldc, + cublasComputeType_t, cublasGemmAlgo_t) { + if (ta == CUBLAS_OP_T && tb == CUBLAS_OP_N && at == CUDA_R_16F && bt == CUDA_R_16F && ct == CUDA_R_32F && lda == k && + strata::dsv41::metal_host::gemm_f16_nt((const uint16_t*) A, (const uint16_t*) B, ldb, (float*) C, ldc, n, m, k, *(const float*) alpha, + *(const float*) beta)) + return CUBLAS_STATUS_SUCCESS; + strata_metal_dense_gemm_missing("cublasGemmEx (this operand combination)"); +} + +template +static inline cublasStatus_t cublasGemmEx_unsupported(A...) { + strata_metal_dense_gemm_missing("cublasGemmEx"); +} diff --git a/edge0-strata/metal/include/metal_compat/cuda_bf16.h b/edge0-strata/metal/include/metal_compat/cuda_bf16.h new file mode 100644 index 00000000..3741b44c --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/cuda_bf16.h @@ -0,0 +1,59 @@ +// metal/include/metal_compat/cuda_bf16.h - the handful of CUDA bf16/fp8 names the dsv41 sources use, on Metal. +// +// Only what the engine touches (found by compiling src/dsv41/*.cu against this directory, 2026-10-23): the storage typedef +// and four conversions. On a Mac the bf16 storage is a uint16_t and the conversion is the standard round-to-nearest-even +// truncation CUDA's __float2bfloat16_rn does; fp8 is E4M3 with the same rounding. +#pragma once + +#include "cuda_fp16.h" // __half_raw, which the fp8 conversions return +#include +#include + +typedef uint16_t __nv_bfloat16; +typedef uint16_t __nv_bfloat162; +typedef uint8_t __nv_fp8_storage_t; + +static inline uint16_t strata_f32_to_bf16_rn(float f) { + uint32_t u; + std::memcpy(&u, &f, sizeof u); + if ((u & 0x7fffffffu) > 0x7f800000u) return (uint16_t) ((u >> 16) | 0x0040u); // NaN keeps a payload + uint32_t r = (u + 0x7fffu + ((u >> 16) & 1u)) >> 16; + return (uint16_t) r; +} + +static inline __nv_bfloat16 __float2bfloat16_rn(float f) { return strata_f32_to_bf16_rn(f); } +static inline uint16_t __bfloat16_as_ushort(__nv_bfloat16 v) { return v; } + +// E4M3 (no infinities, max 448), round to nearest even - the inverse of the half conversion the engine also uses +static inline __nv_fp8_storage_t __nv_cvt_float_to_fp8(float f, int, int) { + if (f != f) return 0x7f; + float a = f < 0 ? -f : f; + int sign = f < 0 ? 0x80 : 0; + if (a == 0.0f) return (__nv_fp8_storage_t) sign; + if (a >= 448.0f) return (__nv_fp8_storage_t) (sign | 0x7e); + int e = 0; + while (a >= 2.0f) { a *= 0.5f; ++e; } + while (a < 1.0f) { a *= 2.0f; --e; } + // a in [1,2): 3 mantissa bits, bias 7 + float scaled = (a - 1.0f) * 8.0f; + int m = (int) (scaled + 0.5f); + int biased = e + 7; + if (m == 8) { m = 0; ++biased; } + if (biased <= 0) return (__nv_fp8_storage_t) sign; // subnormal -> zero here + if (biased >= 15) return (__nv_fp8_storage_t) (sign | 0x7e); // clamp to 448 + return (__nv_fp8_storage_t) (sign | (biased << 3) | m); +} +static inline uint16_t __nv_cvt_fp8_to_halfraw(__nv_fp8_storage_t v, int) { + // half (IEEE binary16) from E4M3: the engine only reads these back for the vision tower's tables + const int sign = (v & 0x80) ? 0x8000 : 0; + const int e = (v >> 3) & 0xf, m = v & 0x7; + if (e == 0 && m == 0) return (uint16_t) sign; + int he, hm; + if (e == 0) { he = 0; hm = m << 7; } + else { he = e - 7 + 15; hm = m << 7; } + return (uint16_t) (sign | (he << 10) | hm); +} + +// the conversion modes __nv_cvt_float_to_fp8 / __nv_cvt_fp8_to_halfraw take (CUDA spells them as an enum) +enum __nv_fp8_interpretation_t { __NV_E4M3 = 0, __NV_E5M2 = 1 }; +enum __nv_saturation_t { __NV_NOSAT = 0, __NV_SATFINITE = 1 }; diff --git a/edge0-strata/metal/include/metal_compat/cuda_device_builtins.h b/edge0-strata/metal/include/metal_compat/cuda_device_builtins.h new file mode 100644 index 00000000..4c990455 --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/cuda_device_builtins.h @@ -0,0 +1,165 @@ +#ifndef __grid_constant__ +#define __grid_constant__ +#endif +#ifndef __device_builtin__ +#define __device_builtin__ +#endif + +// cuda_device_builtins.h - the CUDA device built-ins, declared for the HOST compile only. +// +// The DeepSeek kernels are compiled twice in this port: as C++ here, to enumerate exactly which API the device code uses, +// and as MSL (kernels/*.metal) where the real work happens. These declarations let the first pass run; nothing here is a +// working kernel. Metal spells the same things thread_position_in_grid, threadgroup_barrier, simd_shuffle and so on. +#pragma once + +#include + +struct dim3 { + unsigned x, y, z; + dim3(unsigned _x = 1, unsigned _y = 1, unsigned _z = 1) : x(_x), y(_y), z(_z) {} +}; +struct strata_dim3 { unsigned x, y, z; }; +static strata_dim3 threadIdx{0, 0, 0}, blockIdx{0, 0, 0}, blockDim{1, 1, 1}, gridDim{1, 1, 1}; + +static inline void __syncthreads() {} +static inline void __syncwarp(unsigned = 0xffffffffu) {} +static inline unsigned __ballot_sync(unsigned, int) { return 0; } +static inline int __popc(unsigned v) { return __builtin_popcount(v); } +static inline int __ffs(unsigned v) { return __builtin_ffs((int) v); } + +// the bit and integer intrinsics the block decoders use. Each is the exact CUDA semantics, computed on the host here; the +// MSL kernels implement the same operations with Metal's own select/rotate/permute built-ins. +static inline unsigned __byte_perm(unsigned x, unsigned y, unsigned s) { + unsigned out = 0; + const unsigned long long src = ((unsigned long long) y << 32) | x; + for (int i = 0; i < 4; ++i) { + const unsigned sel = (s >> (4 * i)) & 0x7u; + out |= (unsigned) ((src >> (8 * sel)) & 0xffu) << (8 * i); + } + return out; +} +static inline int __mul24(int a, int b) { return a * b; } +static inline unsigned __umul24(unsigned a, unsigned b) { return a * b; } +static inline unsigned __brev(unsigned v) { return __builtin_bitreverse32(v); } +static inline int __clz(int v) { return v ? __builtin_clz((unsigned) v) : 32; } +static inline int __clzll(long long v) { return v ? __builtin_clzll((unsigned long long) v) : 64; } +static inline float __fdividef(float a, float b) { return a / b; } +static inline float __saturatef(float a) { return a < 0.0f ? 0.0f : (a > 1.0f ? 1.0f : a); } +static inline float __int_as_float(int v) { float f; __builtin_memcpy(&f, &v, 4); return f; } +static inline int __float_as_int(float f) { int v; __builtin_memcpy(&v, &f, 4); return v; } +static inline unsigned __float_as_uint(float f) { unsigned v; __builtin_memcpy(&v, &f, 4); return v; } +static inline float __uint_as_float(unsigned v) { float f; __builtin_memcpy(&f, &v, 4); return f; } +static inline void __threadfence() {} +static inline void __threadfence_block() {} +static inline float atomicAdd(float* p, float v) { const float o = *p; *p = o + v; return o; } +static inline int atomicAdd(int* p, int v) { const int o = *p; *p = o + v; return o; } +static inline unsigned atomicAdd(unsigned* p, unsigned v) { const unsigned o = *p; *p = o + v; return o; } +static inline int atomicMax(int* p, int v) { const int o = *p; if (v > o) *p = v; return o; } +static inline int atomicCAS(int* p, int c, int v) { const int o = *p; if (o == c) *p = v; return o; } + +// the SIMD video instructions (packed 8-bit lanes) the i-quant block decoders use: on CUDA these are one instruction, on +// Apple Silicon NEON has them directly (vsubq_u8, vpaddlq / vdotq_s32) and the MSL kernels use those. Here they are the +// exact per-lane semantics, so the host compile reproduces the same numbers. +static inline unsigned __vsub4(unsigned a, unsigned b) { + unsigned out = 0; + for (int i = 0; i < 4; ++i) { + const int d = (int) ((a >> (8 * i)) & 0xffu) - (int) ((b >> (8 * i)) & 0xffu); + out |= (unsigned) ((d < 0 ? 0 : d) & 0xff) << (8 * i); // unsigned saturated + } + return out; +} +static inline unsigned __vsubss4(unsigned a, unsigned b) { + unsigned out = 0; + for (int i = 0; i < 4; ++i) { + int d = (int) ((a >> (8 * i)) & 0xffu) - (int) ((b >> (8 * i)) & 0xffu); + if (d < -128) d = -128; + if (d > 127) d = 127; + out |= (unsigned) (d & 0xff) << (8 * i); + } + return out; +} +static inline unsigned __vadd4(unsigned a, unsigned b) { + unsigned out = 0; + for (int i = 0; i < 4; ++i) { + const int d = (int) ((a >> (8 * i)) & 0xffu) + (int) ((b >> (8 * i)) & 0xffu); + out |= (unsigned) ((d > 255 ? 255 : d) & 0xff) << (8 * i); + } + return out; +} +static inline unsigned __vaddss4(unsigned a, unsigned b) { + unsigned out = 0; + for (int i = 0; i < 4; ++i) { + int d = (int) ((a >> (8 * i)) & 0xffu) + (int) ((b >> (8 * i)) & 0xffu); + if (d > 127) d = 127; + out |= (unsigned) (d & 0xff) << (8 * i); + } + return out; +} +static inline unsigned __vcmpeq4(unsigned a, unsigned b) { + unsigned out = 0; + for (int i = 0; i < 4; ++i) { + const unsigned ea = (a >> (8 * i)) & 0xffu, eb = (b >> (8 * i)) & 0xffu; + out |= (ea == eb ? 0xffu : 0u) << (8 * i); + } + return out; +} +static inline unsigned __vcmpgtu4(unsigned a, unsigned b) { + unsigned out = 0; + for (int i = 0; i < 4; ++i) { + const unsigned ea = (a >> (8 * i)) & 0xffu, eb = (b >> (8 * i)) & 0xffu; + out |= (ea > eb ? 0xffu : 0u) << (8 * i); + } + return out; +} +static inline int __dp4a(unsigned a, unsigned b, int c) { + for (int i = 0; i < 4; ++i) c += (int) (signed char) ((a >> (8 * i)) & 0xffu) * (int) (signed char) ((b >> (8 * i)) & 0xffu); + return c; +} + +// the remaining scalar math helpers CUDA provides and the decoders use (round-to-nearest, plus min/max for integers - CUDA +// brings those in from its own headers, the C++ standard library does not for every type) +static inline float __fsub_rn(float a, float b) { return a - b; } +static inline float __fdiv_rn(float a, float b) { return a / b; } +static inline float __frcp_rn_(float a) { return 1.0f / a; } +template static inline T strata_min(T a, T b) { return a < b ? a : b; } +template static inline T strata_max(T a, T b) { return a > b ? a : b; } + +// the last names the kernel sources use (a full-list compile found twelve families, 2026-10-23): the grid type above, a +// system fence, the funnel shift, the hardware transcendentals and the read-only/volatile loads. +static inline void __threadfence_system() {} +static inline unsigned __funnelshift_r(unsigned lo, unsigned hi, unsigned shift) { + const unsigned long long v = ((unsigned long long) hi << 32) | lo; + return (unsigned) ((v >> (shift & 31)) & 0xffffffffu); +} +static inline unsigned __funnelshift_l(unsigned lo, unsigned hi, unsigned shift) { + const unsigned long long v = ((unsigned long long) hi << 32) | lo; + return (unsigned) (((v << (shift & 31)) >> 32) & 0xffffffffu); +} +static inline float rsqrtf(float a) { return 1.0f / __builtin_sqrtf(a); } +static inline float __expf(float a) { return __builtin_expf(a); } +static inline float __logf(float a) { return __builtin_logf(a); } +static inline float __sinf(float a) { return __builtin_sinf(a); } +static inline float __cosf(float a) { return __builtin_cosf(a); } +static inline void sincosf(float a, float* s, float* c) { *s = __builtin_sinf(a); *c = __builtin_cosf(a); } +static inline int __float2int_rn(float a) { return (int) __builtin_roundf(a); } +static inline int __float2int_rz(float a) { return (int) a; } +static inline int __float2int_rd(float a) { return (int) __builtin_floorf(a); } +static inline int __float2int_ru(float a) { return (int) __builtin_ceilf(a); } +template static inline T __ldg(const T* p) { return *p; } +template static inline T __ldcv(const T* p) { return *p; } + +// shuffles: one template each, so int/unsigned/float/short all bind (CUDA has overloads; a template is the same thing here) +template static inline T __shfl_xor_sync(unsigned, T v, int) { return v; } +template static inline T __shfl_down_sync(unsigned, T v, int) { return v; } +template static inline T __shfl_up_sync(unsigned, T v, int) { return v; } +template static inline T __shfl_sync(unsigned, T v, int) { return v; } + +// the integer multiply-high and the bitwise atomics the sampler uses +static inline unsigned __umulhi(unsigned a, unsigned b) { return (unsigned) (((unsigned long long) a * b) >> 32); } +static inline int __mulhi(int a, int b) { return (int) (((long long) a * b) >> 32); } +static inline unsigned atomicOr(unsigned* p, unsigned v) { const unsigned o = *p; *p = o | v; return o; } +static inline unsigned atomicAnd(unsigned* p, unsigned v) { const unsigned o = *p; *p = o & v; return o; } +static inline unsigned atomicXor(unsigned* p, unsigned v) { const unsigned o = *p; *p = o ^ v; return o; } +static inline unsigned atomicExch(unsigned* p, unsigned v) { const unsigned o = *p; *p = v; return o; } +static inline unsigned atomicInc(unsigned* p, unsigned v) { const unsigned o = *p; *p = (o >= v ? 0 : o + 1); return o; } +static inline unsigned atomicDec(unsigned* p, unsigned v) { const unsigned o = *p; *p = (o == 0 || o > v ? v : o - 1); return o; } diff --git a/edge0-strata/metal/include/metal_compat/cuda_fp16.h b/edge0-strata/metal/include/metal_compat/cuda_fp16.h new file mode 100644 index 00000000..1f88d298 --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/cuda_fp16.h @@ -0,0 +1,99 @@ +// metal/include/metal_compat/cuda_fp16.h - the half-precision surface the dsv41 sources use, on a Mac. +// +// The engine touches exactly these (counted by compiling src/dsv41/*.cu against this directory, M5 Max 2026-10-23): +// __half 16, __half2 11, __half_raw 3, __half_as_ushort 8, +// __float2half_rn 17, __half2float 18, __half22float2 6, +// the round-to-nearest float helpers (__fmul_rn, __fadd_rn, __fmaf_rn, __frcp_rn, __fsqrt_rn), +// and the plain half/half2 typedefs ggml-common.h uses. +// On Apple Silicon half is IEEE binary16 too, so the storage is a uint16_t and the conversions are the standard +// round-to-nearest-even ones - the same bits CUDA produces. +#pragma once + +#include "vector_types.h" + +#include +#include + +struct __half_raw { + uint16_t x; + __half_raw() = default; + constexpr __half_raw(uint16_t bits) : x(bits) {} +}; + +struct __half { + uint16_t x; + __half() = default; + __half(const __half_raw& r) : x(r.x) {} + operator __half_raw() const { return __half_raw{x}; } +}; + +struct __half2 { + uint16_t x, y; +}; + +typedef __half half; +typedef __half2 half2; + +static inline uint16_t strata_f32_to_f16_bits(float f) { + uint32_t u; + std::memcpy(&u, &f, sizeof u); + const uint32_t sign = (u >> 16) & 0x8000u; + const uint32_t mag = u & 0x7fffffffu; + if (mag >= 0x7f800000u) return (uint16_t) (sign | (mag > 0x7f800000u ? 0x7e00u : 0x7c00u)); + if (mag >= 0x477ff000u) return (uint16_t) (sign | 0x7c00u); + const uint32_t r = (mag + 0x0fffu + ((mag >> 13) & 1u)) >> 13; + if (r >= 0x7c00u) return (uint16_t) (sign | 0x7c00u); + return (uint16_t) (sign | r); +} + +static inline float strata_f16_bits_to_f32(uint16_t h) { + const uint32_t sign = (uint32_t) (h & 0x8000u) << 16; + const uint32_t e = (h >> 10) & 0x1fu, m = h & 0x3ffu; + uint32_t bits; + if (e == 0) { + if (m == 0) bits = sign; + else { + const float v = (float) m * (1.0f / 16777216.0f); + return sign ? -v : v; + } + } else if (e == 31) { + bits = sign | 0x7f800000u | (m << 13); + } else { + bits = sign | ((e - 15 + 127) << 23) | (m << 13); + } + float out; + std::memcpy(&out, &bits, sizeof out); + return out; +} + +static inline __half __float2half_rn(float f) { __half h; h.x = strata_f32_to_f16_bits(f); return h; } +static inline __half __float2half(float f) { return __float2half_rn(f); } +static inline float __half2float(__half h) { return strata_f16_bits_to_f32(h.x); } +static inline uint16_t __half_as_ushort(__half h) { return h.x; } +static inline __half __ushort_as_half(uint16_t u) { __half h; h.x = u; return h; } +static inline float2 __half22float2(__half2 h) { return make_float2(strata_f16_bits_to_f32(h.x), strata_f16_bits_to_f32(h.y)); } +static inline __half2 __floats2half2_rn(float a, float b) { + __half2 h; + h.x = strata_f32_to_f16_bits(a); + h.y = strata_f32_to_f16_bits(b); + return h; +} + +// round-to-nearest float helpers: the FP ops here are IEEE and un-fused, which is what these mean +static inline float __fmul_rn(float a, float b) { return a * b; } +static inline float __fadd_rn(float a, float b) { return a + b; } +static inline float __fmaf_rn(float a, float b, float c) { return a * b + c; } +static inline float __frcp_rn(float a) { return 1.0f / a; } +static inline float __fsqrt_rn(float a) { return __builtin_sqrtf(a); } + +// the half2 pack/unpack and arithmetic helpers the trellis and vision paths use +static inline __half2 __halves2half2(__half lo, __half hi) { __half2 h; h.x = lo.x; h.y = hi.x; return h; } +static inline __half2 __half2half2(__half v) { __half2 h; h.x = v.x; h.y = v.x; return h; } +static inline __half __low2half(__half2 h) { __half o; o.x = h.x; return o; } +static inline __half __high2half(__half2 h) { __half o; o.x = h.y; return o; } +static inline float __low2float(__half2 h) { return strata_f16_bits_to_f32(h.x); } +static inline float __high2float(__half2 h) { return strata_f16_bits_to_f32(h.y); } +static inline __half2 __hadd2(__half2 a, __half2 b) { return __floats2half2_rn(__half2float(__low2half(a)) + __half2float(__low2half(b)), __half2float(__high2half(a)) + __half2float(__high2half(b))); } +static inline __half2 __hmul2(__half2 a, __half2 b) { return __floats2half2_rn(__half2float(__low2half(a)) * __half2float(__low2half(b)), __half2float(__high2half(a)) * __half2float(__high2half(b))); } +static inline __half2 __hfma2(__half2 a, __half2 b, __half2 c) { return __floats2half2_rn(__half2float(__low2half(a)) * __half2float(__low2half(b)) + __half2float(__low2half(c)), __half2float(__high2half(a)) * __half2float(__high2half(b)) + __half2float(__high2half(c))); } +static inline __half2 __hneg2(__half2 a) { __half2 h; h.x = a.x ^ 0x8000u; h.y = a.y ^ 0x8000u; return h; } diff --git a/edge0-strata/metal/include/metal_compat/cuda_fp8.h b/edge0-strata/metal/include/metal_compat/cuda_fp8.h new file mode 100644 index 00000000..c0f63abb --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/cuda_fp8.h @@ -0,0 +1,4 @@ +// metal/include/metal_compat/cuda_fp8.h - the fp8 names the engine includes; the definitions live beside the bf16 ones. +#pragma once + +#include "cuda_bf16.h" diff --git a/edge0-strata/metal/include/metal_compat/cuda_launch.h b/edge0-strata/metal/include/metal_compat/cuda_launch.h new file mode 100644 index 00000000..dc0f2bb3 --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/cuda_launch.h @@ -0,0 +1,31 @@ +// cuda_launch.h - the <<>> launch syntax, as a call. +// +// A C++ compiler cannot parse CUDA's launch syntax, and the Metal port replaces every launch with a dispatch anyway (the grid +// becomes a threadgroup count, the block a threadgroup size). strata_ds4_launch() is that call: the kernel first, then the +// launch shape, then the kernel's own arguments - every parameter generic, because the shapes vary (dim3, plain ints, sizes, +// and the stream type), and the point of this header is to make the sources parse, not to type-check the port. +// +// Until a kernel has a Metal implementation the call stops the run naming the kernel, so the engine reports the first one it +// needs and the port proceeds one kernel at a time. +#pragma once + +#include +#include + +template +static inline void strata_ds4_launch(K kernel, A&&...) { + static_assert(sizeof(K) > 0, "kernel"); + std::fprintf(stderr, "strata metal: a dsv41 kernel launch has no Metal dispatch yet\n"); + std::abort(); +} + +// the call sites sit inside namespaces whose first argument would win argument-dependent lookup: force the global one +// the call sites carry their own location, so a run names the exact launch that needs a dispatch +#define strata_ds4_launch(...) ::strata_ds4_launch_at(__FILE__, __LINE__, __VA_ARGS__) + +// the macro adds the call site's location, so a run names the exact launch that still needs a Metal dispatch +template +static inline void strata_ds4_launch_at(const char* file, int line, K, A&&...) { + std::fprintf(stderr, "strata metal: kernel launch with no Metal dispatch at %s:%d\n", file, line); + std::abort(); +} diff --git a/edge0-strata/metal/include/metal_compat/cuda_runtime.h b/edge0-strata/metal/include/metal_compat/cuda_runtime.h new file mode 100644 index 00000000..29aee70a --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/cuda_runtime.h @@ -0,0 +1,218 @@ +// The CUDA function-space qualifiers and the few launch/kernel spellings the shared sources use. On Metal they say +#include "cuda_launch.h" +// nothing: a "kernel" here is an MSL function compiled at run time and dispatched by the host runtime, so the qualifiers +// are empty and __forceinline__ is plain inline. Part of the metal_compat layer (M5 Max, 2026-10-23). +#ifndef __host__ +#define __host__ +#endif +#ifndef __device__ +#define __device__ +#endif +#ifndef __global__ +#define __global__ +#endif +#ifndef __shared__ +#define __shared__ +#endif +#ifndef __constant__ +#define __constant__ +#endif +#ifndef __forceinline__ +#define __forceinline__ inline +#endif +#ifndef __launch_bounds__ +#define __launch_bounds__(...) +#endif +#ifndef __restrict__ +#define __restrict__ +#endif + +// metal/include/metal_compat/cuda_runtime.h - the CUDA runtime names the host files use, on Metal. +// +// The HIP build's approach (include/strata/hip_compat/cuda_runtime.h maps every cudaX the host code calls to hipX), +// applied to Metal: this directory goes first on the Metal build's include path, so the shared host files compile for +// Metal UNCHANGED, and the CUDA / HIP builds never see this file (their objects stay byte-identical by construction). +// Every name forwards to strata::gpu's Metal implementation (metal/src/gpu/runtime.cpp), whose semantics are the +// CUDA ones. Only the names in the Phase 1 inventory (phaseA-tests/mac-design.md §1.3) are provided; a host file that +// uses another one fails to compile on Metal, which is the point. +#pragma once + +// CUDA's alignment qualifier +#ifndef __align__ +#define __align__(n) alignas(n) +#endif + +#include "cuda_device_builtins.h" + +#include + +using std::fmaxf; +using std::fminf; +using std::sqrtf; +using std::fabsf; +using std::powf; +using std::expf; +using std::logf; + +// CUDA's headers put min/max in scope for integers; the sources rely on it +template static inline T min(T a, T b) { return strata_min(a, b); } +template static inline T max(T a, T b) { return strata_max(a, b); } + +#include +#include + +#include // metal_compat/immintrin.h: CUDA's headers bring the x86 one in on a PC +#include "strata/gpu/runtime.hpp" // metal/include/strata/gpu: the Metal runtime + +namespace strata::gpu::compat_detail { inline Error bad() { return kErrorUnsupported; } } + +// ---- types and constants +using cudaError_t = strata::gpu::Error; +using cudaStream_t = strata::gpu::Stream; +using cudaEvent_t = strata::gpu::Event; +using cudaGraph_t = strata::gpu::Graph; +using cudaGraphExec_t = strata::gpu::GraphExec; +using cudaGraphNode_t = void*; +using cudaDeviceProp = strata::gpu::DeviceProp; +using cudaMemcpyKind = strata::gpu::CopyKind; +using cudaHostFn_t = strata::gpu::HostFn; +inline constexpr cudaError_t cudaSuccess = strata::gpu::kSuccess; +inline constexpr cudaError_t cudaErrorNotReady = strata::gpu::kNotReady; +inline constexpr cudaError_t cudaErrorMemoryAllocation = strata::gpu::kErrorMemoryAllocation; +inline constexpr cudaError_t cudaErrorPeerAccessAlreadyEnabled = strata::gpu::kErrorPeerAccessAlreadyEnabled; +inline constexpr cudaError_t cudaErrorStreamCaptureUnsupported = strata::gpu::kErrorUnsupported; +inline constexpr cudaMemcpyKind cudaMemcpyHostToDevice = strata::gpu::kHostToDevice; +inline constexpr cudaMemcpyKind cudaMemcpyDeviceToHost = strata::gpu::kDeviceToHost; +inline constexpr cudaMemcpyKind cudaMemcpyDeviceToDevice = strata::gpu::kDeviceToDevice; +inline constexpr cudaMemcpyKind cudaMemcpyDefault = strata::gpu::kCopyDefault; +inline constexpr unsigned cudaHostAllocDefault = strata::gpu::kHostAllocDefault; +inline constexpr unsigned cudaHostAllocPortable = strata::gpu::kHostAllocPortable; +inline constexpr unsigned cudaHostAllocMapped = strata::gpu::kHostAllocMapped; +inline constexpr unsigned cudaHostRegisterPortable = strata::gpu::kHostRegisterPortable; +inline constexpr unsigned cudaHostRegisterMapped = strata::gpu::kHostRegisterMapped; +inline constexpr unsigned cudaHostRegisterReadOnly = strata::gpu::kHostRegisterReadOnly; +inline constexpr unsigned cudaEventDefault = strata::gpu::kEventDefault; +inline constexpr unsigned cudaEventDisableTiming = strata::gpu::kEventDisableTiming; +inline constexpr unsigned cudaEventBlockingSync = strata::gpu::kEventBlockingSync; +inline constexpr unsigned cudaStreamNonBlocking = strata::gpu::kStreamNonBlocking; +enum cudaStreamCaptureStatus { cudaStreamCaptureStatusNone = 0, cudaStreamCaptureStatusActive = 1, cudaStreamCaptureStatusInvalidated = 2 }; +static inline strata::gpu::Error cudaStreamIsCapturing(strata::gpu::Stream, cudaStreamCaptureStatus* st) { + if (st) *st = cudaStreamCaptureStatusNone; + return strata::gpu::kSuccess; +} + +enum cudaStreamCaptureMode { cudaStreamCaptureModeGlobal = 0, cudaStreamCaptureModeThreadLocal = 1, cudaStreamCaptureModeRelaxed = 2 }; +// (the compute capability is asked by the 0.1.45 Turing defaults in verify.cpp: no answer here, which reads as "not a Turing card") +enum cudaDeviceAttr { + cudaDevAttrClockRate = 13, cudaDevAttrMultiProcessorCount = 16, cudaDevAttrIntegrated = 18, + cudaDevAttrComputeCapabilityMajor = 75, cudaDevAttrComputeCapabilityMinor = 76 +}; + +// ---- errors, devices +inline cudaError_t cudaGetLastError() { return strata::gpu::get_last_error(); } +inline cudaError_t cudaPeekAtLastError() { return strata::gpu::peek_last_error(); } +inline const char* cudaGetErrorString(cudaError_t e) { return strata::gpu::error_string(e); } +inline cudaError_t cudaGetDevice(int* d) { return strata::gpu::get_device(d); } +inline cudaError_t cudaSetDevice(int d) { return strata::gpu::set_device(d); } +inline cudaError_t cudaGetDeviceCount(int* n) { return strata::gpu::device_count(n); } +inline cudaError_t cudaDeviceSynchronize() { return strata::gpu::device_synchronize(); } +inline cudaError_t cudaMemGetInfo(size_t* f, size_t* t) { return strata::gpu::mem_get_info(f, t); } +inline cudaError_t cudaGetDeviceProperties(cudaDeviceProp* p, int d) { return strata::gpu::get_device_properties(p, d); } +inline cudaError_t cudaRuntimeGetVersion(int* v) { return strata::gpu::runtime_version(v); } +inline cudaError_t cudaDriverGetVersion(int* v) { return strata::gpu::runtime_version(v); } +inline cudaError_t cudaDeviceGetAttribute(int* v, cudaDeviceAttr a, int d) { + switch (a) { + case cudaDevAttrIntegrated: return strata::gpu::device_integrated(v, d); + case cudaDevAttrMultiProcessorCount: return strata::gpu::device_sm_count(v, d); + case cudaDevAttrClockRate: return strata::gpu::device_clock_khz(v, d); + case cudaDevAttrComputeCapabilityMajor: + case cudaDevAttrComputeCapabilityMinor: return strata::gpu::compat_detail::bad(); + } + return strata::gpu::compat_detail::bad(); +} +inline cudaError_t cudaDeviceCanAccessPeer(int* v, int d, int p) { return strata::gpu::device_can_access_peer(v, d, p); } +inline cudaError_t cudaDeviceEnablePeerAccess(int p, unsigned f) { return strata::gpu::device_enable_peer_access(p, f); } +inline cudaError_t cudaDeviceDisablePeerAccess(int) { return cudaSuccess; } // a Mac has one GPU: nothing to disable + +// ---- memory and copies +// the deferred command buffer mtl_rt batches the port's dispatches into: a copy (or any host read of a device result) +// has to commit it first, or it reads memory the GPU has not written yet. Defined in ds4_finish.cpp. +extern "C" void ds4_metal_flush(); +extern "C" void ds4_metal_kick(); + +template inline cudaError_t cudaMalloc(T** p, size_t n) { return strata::gpu::malloc_device(p, n); } +inline cudaError_t cudaFree(void* p) { return strata::gpu::free_device(p); } +template inline cudaError_t cudaHostAlloc(T** p, size_t n, unsigned f) { return strata::gpu::host_alloc(p, n, f); } +template inline cudaError_t cudaMallocHost(T** p, size_t n) { return strata::gpu::malloc_host(p, n); } +inline cudaError_t cudaFreeHost(void* p) { return strata::gpu::free_host(p); } +template inline cudaError_t cudaHostGetDevicePointer(T** d, void* h, unsigned f) { return strata::gpu::host_device_pointer(d, h, f); } +inline cudaError_t cudaHostRegister(void* p, size_t n, unsigned f) { return strata::gpu::host_register(p, n, f); } +inline cudaError_t cudaHostUnregister(void* p) { return strata::gpu::host_unregister(p); } +inline cudaError_t cudaMemset(void* p, int v, size_t n) { return strata::gpu::memset_device(p, v, n); } +inline cudaError_t cudaMemsetAsync(void* p, int v, size_t n, cudaStream_t s = nullptr) { return strata::gpu::memset_async(p, v, n, s); } +inline cudaError_t cudaMemcpy(void* d, const void* s, size_t n, cudaMemcpyKind k) { return strata::gpu::memcpy_sync(d, s, n, k); } +inline cudaError_t cudaMemcpyAsync(void* d, const void* s, size_t n, cudaMemcpyKind k, cudaStream_t st = nullptr) { + return strata::gpu::memcpy_async(d, s, n, k, st); +} +inline cudaError_t cudaMemcpy2DAsync(void* d, size_t dp, const void* s, size_t sp, size_t w, size_t h, cudaMemcpyKind k, + cudaStream_t st = nullptr) { + return strata::gpu::memcpy_2d_async(d, dp, s, sp, w, h, k, st); +} +inline cudaError_t cudaMemcpyPeerAsync(void* d, int dd, const void* s, int sd, size_t n, cudaStream_t st = nullptr) { + return strata::gpu::memcpy_peer_async(d, dd, s, sd, n, st); +} +inline cudaError_t cudaMemcpyPeer(void* d, int dd, const void* s, int sd, size_t n) { + const cudaError_t e = cudaMemcpyPeerAsync(d, dd, s, sd, n, nullptr); + return e != cudaSuccess ? e : cudaDeviceSynchronize(); +} + +// ---- streams, events, host functions +inline cudaError_t cudaStreamCreate(cudaStream_t* s) { return strata::gpu::stream_create(s); } +inline cudaError_t cudaStreamCreateWithFlags(cudaStream_t* s, unsigned f) { return strata::gpu::stream_create(s, f); } +inline cudaError_t cudaDeviceGetStreamPriorityRange(int*, int*) { return strata::gpu::compat_detail::bad(); } // no priorities: callers fall back to a plain stream +inline cudaError_t cudaStreamCreateWithPriority(cudaStream_t* s, unsigned f, int) { return strata::gpu::stream_create(s, f); } +inline cudaError_t cudaStreamDestroy(cudaStream_t s) { return strata::gpu::stream_destroy(s); } +inline cudaError_t cudaStreamSynchronize(cudaStream_t s) { return strata::gpu::stream_synchronize(s); } +inline cudaError_t cudaStreamQuery(cudaStream_t s) { return strata::gpu::stream_query(s); } +inline cudaError_t cudaStreamWaitEvent(cudaStream_t s, cudaEvent_t e, unsigned f = 0) { return strata::gpu::stream_wait_event(s, e, f); } +inline cudaError_t cudaLaunchHostFunc(cudaStream_t s, cudaHostFn_t fn, void* a) { return strata::gpu::launch_host_func(s, fn, a); } +inline cudaError_t cudaEventCreate(cudaEvent_t* e) { return strata::gpu::event_create(e); } +inline cudaError_t cudaEventCreateWithFlags(cudaEvent_t* e, unsigned f) { return strata::gpu::event_create(e, f); } +inline cudaError_t cudaEventDestroy(cudaEvent_t e) { return strata::gpu::event_destroy(e); } +inline cudaError_t cudaEventRecord(cudaEvent_t e, cudaStream_t s = nullptr) { return strata::gpu::event_record(e, s); } +inline cudaError_t cudaEventQuery(cudaEvent_t e) { return strata::gpu::event_query(e); } +inline cudaError_t cudaEventSynchronize(cudaEvent_t e) { return strata::gpu::event_synchronize(e); } +inline cudaError_t cudaEventElapsedTime(float* ms, cudaEvent_t a, cudaEvent_t b) { return strata::gpu::event_elapsed_ms(ms, a, b); } + +// ---- graphs +inline cudaError_t cudaStreamBeginCapture(cudaStream_t s, cudaStreamCaptureMode m) { + return m == cudaStreamCaptureModeThreadLocal ? strata::gpu::begin_capture(s) : strata::gpu::compat_detail::bad(); +} +inline cudaError_t cudaStreamEndCapture(cudaStream_t s, cudaGraph_t* g) { return strata::gpu::end_capture(s, g); } +inline cudaError_t cudaGraphGetNodes(cudaGraph_t g, cudaGraphNode_t* nodes, size_t* n) { + return nodes == nullptr ? strata::gpu::graph_node_count(g, n) : strata::gpu::compat_detail::bad(); +} +inline cudaError_t cudaGraphInstantiate(cudaGraphExec_t* x, cudaGraph_t g, cudaGraphNode_t*, char*, size_t) { + return strata::gpu::graph_instantiate(x, g); +} +inline cudaError_t cudaGraphInstantiate(cudaGraphExec_t* x, cudaGraph_t g, unsigned long long f = 0) { + return strata::gpu::graph_instantiate(x, g, f); +} +inline cudaError_t cudaGraphLaunch(cudaGraphExec_t x, cudaStream_t s) { return strata::gpu::graph_launch(x, s); } +inline cudaError_t cudaGraphUpload(cudaGraphExec_t x, cudaStream_t s) { return strata::gpu::graph_upload(x, s); } +inline cudaError_t cudaGraphDestroy(cudaGraph_t g) { return strata::gpu::graph_destroy(g); } +inline cudaError_t cudaGraphExecDestroy(cudaGraphExec_t x) { return strata::gpu::graph_exec_destroy(x); } + +// cudaMemcpy2D(dst, dpitch, src, spitch, width, height, kind): a strided copy, used by the vision tower's patch rows and +// the Engram tables' column slices. The Metal runtime has no pitched copy yet, so this reports failure rather than +// copying nothing silently. +inline strata::gpu::Error cudaMemcpy2D(void*, size_t, const void*, size_t, size_t, size_t, strata::gpu::CopyKind) { + return strata::gpu::kErrorUnsupported; // cudaMemcpy2D_unimplemented_marker +} + +// occupancy query: the Metal runtime decides its own threadgroup limits, so answer "one block" and let the caller size grids +template +static inline strata::gpu::Error cudaOccupancyMaxActiveBlocksPerMultiprocessor(int* n, K&&, int, A&&...) { + if (n) *n = 1; + return strata::gpu::kSuccess; +} diff --git a/edge0-strata/metal/include/metal_compat/ds4_metal_host.h b/edge0-strata/metal/include/metal_compat/ds4_metal_host.h new file mode 100644 index 00000000..9eb7c9a5 --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/ds4_metal_host.h @@ -0,0 +1,111 @@ +// metal/include/metal_compat/ds4_metal_host.h - the dsv41 kernels that have a Metal dispatch. +// +// GENERATED from metal/src/engine/ds4_host.cpp and ds4_finish.cpp, which are the source of truth: do not hand-edit the +// declaration block, regenerate it (see RIGS/m1002.md). +#pragma once + +#include +#include +#include + +namespace strata::gpu { +struct StreamImpl; +using Stream = StreamImpl*; +} // namespace strata::gpu + +namespace strata::dsv41::metal_host { + +// ds4_kernels.cu:1894 finish_kernel - the engine's Finish, as ds4_kernels.cuh declares it +struct FinishIn { + const float* w; + float eps; + int rd, pos, inverse, vpp, pos_step; + const float* freqs; + int quant; + const uint16_t* dst; + long long dst_stride; + uint8_t* pdst; + long long pdst_stride; +}; + +bool embed(const uint16_t* rows, int T, float* h, int dim, int hc, void*); +bool engram_dequant_tok(const uint8_t* src, size_t per, int cols, int T, int hd, float* out, int fmt); +bool ext_to_ring(const void* ext, void* ring, long long p0, int n, int T, int d, size_t row_bytes, void* backup); +bool finish(float* x, int n_vec, int len, long long stride, const FinishIn& f); +bool gemv_f32(const float* W, const float* x, int n, int k, int T, float* y, int ldy); +bool hc_mix(const float* h, const float* fn, int hcdim, int n_mix, int T, float eps, float* mixes); +bool hc_mix_scale(const float* h, int hcdim, int n_mix, int T, float eps, float* mixes); +bool hc_post(const float* y, float* h, const float* post, const float* comb, int dim, int hc, int T); +bool hc_post_v(const float* y, const float* res, const float* post, const float* comb, int dim, int hc, int T, float* out); +bool hc_pre_norm(const float* h, const float* mix, const float* w, int dim, int hc, int T, float eps, float* x); +bool hc_split(const float* mixes, const float* scale, const float* base, int hc, int iters, float hc_eps, int T, float* pre, float* post, float* comb); +bool hc_mix_split(const float* h, const float* fn, int hcdim, int n_mix, int T, float eps, float* mixes, const float* scale, const float* base, int hc, int iters, float hc_eps, float* pre, float* post, float* comb); +bool init_pre(float* pre, int T, void*); +bool iq_mv(const void* jobs, int nj); +bool moe_combine(void* acc, const void* y, const int* pos, const float* w_, int T, int k, int r0, int r1, int dim); +bool moe_sort(const int* idx, int T, int k, const int* vs, int n_vs, int* bounds, int* row_tok, int* pos); +bool moe_win_group_cpu(const void* plan_ptr, void* grp_ptr, void* grp_start, void* grp_ent, void* n_groups, int T, int k); +bool mx_mv(const void* jobs, int nj, int max_n, size_t job_bytes, int fmt); +bool q8_quant_rows(const float* x, size_t ldx, int rows, int k, int8_t* q, float* d); +bool ring_to_ext(const void* ring, void* ext, long long p0, int n, int d, size_t row_bytes); +bool route_publish(const float* logits, const float* bias, int n, int k, int T, float route_scale, const void* vram, int* idx, float* w, void* plan_ptr, const float* bias_vl, unsigned vl); +bool sgemm_tn(const float* A, const float* B, float* C, int m, int nt, int k, int lda, int ldb, int ldc); +bool rows_f16(const float* x, int ldx, const int* idx, int nrows, int ncols, uint16_t* out); +uint64_t zc_map(const void* p, size_t bytes); // zero-copy weights: the GPU address of host memory, 0 if not mappable +bool q8_quant_multi(const void* specs, int n); // {const float* x; uint64 ldx; uint32 rows, k; int8* q; float* d} each +bool iq_set_gu(const uint8_t* wbase, size_t stride, const size_t* tr, const int* hbounds, int n_exp, int total_rows, int n_embd, int n_ff, float limit, const float* xf, int ldx, const int* idx, uint16_t* h, int ldh, int tg, int td); +bool iq_set_dn(const uint8_t* wbase, size_t stride, const size_t* tr, const int* hbounds, int n_exp, int total_rows, int n_embd, int n_ff, float limit, const uint16_t* h16, int ldh, float* y, int ldy, int tg, int td); +bool dequant_rows(int kind, const void* raw, size_t ld, int k, int n, uint16_t* out); +bool gemm_f16_nt(const uint16_t* W16, const uint16_t* X16, int ldx, float* Y, int ldy, int rows, int N, int K, float alpha, float beta); +bool swiglu(const float* g, const float* u, int n, float lim, float* h); +bool unpack_sel(const void* ckv, const int* sel, int sel_stride, int gen, long long p0, int t0, int R, int kcap, const int* nsel, int cap, void* out, int T, int nsel_count, size_t row_bytes); + +// generated from ds4_host.cpp and ds4_finish.cpp +bool block_max(const float* score, long long stride, int T, const void* rows, int blk, float* bm, long long bm_stride, int nb_max); +bool combine(const float* ed, const float* cpu, const float* sh, const void* c, int T, int dim, float* out); +bool compress_win(const float* kv, const float* score, float* kv_state, float* score_state, int ratio, int p0, int T, int d, const float* w, float eps, float* latent); +bool confidence(const float* x, const uint16_t* E, const int* tok, const float* w, int T, int dim, int R, float* out); +bool embed(const uint16_t* rows, int T, float* h, int dim, int hc, void*); +bool embed_addr(const uint16_t* rows, int T, float* h, int dim, int hc); +bool engram_dequant(const uint8_t* rows, const uint8_t* scales, int n, int hd, float* out, int fmt); +bool engram_dequant_tok(const uint8_t* src, size_t per, int cols, int T, int hd, float* out, int fmt); +bool engram_gate(float* h, const float* kv, const float* w, int hc, int dim, int T, float eps); +bool ext_to_ring(const void* ext, void* ring, long long p0, int n, int T, int d, size_t row_bytes, void* backup); +bool finish(float* x, int n_vec, int len, long long stride, const FinishIn& f); +bool gather_rows(const float* x, const void* g, int n, int cols, float* out); +bool gemv_f16_rows(const uint16_t* W, const float* x, size_t ldx, int n, int k, int T, float alpha, float* y, size_t ldy); +bool gemv_f32(const float* W, const float* x, int n, int k, int T, float* y, int ldy); +bool hc_mean(const float* h, int T, int hc, int dim, float* out, long long out_stride); +bool hc_mix(const float* h, const float* fn, int hcdim, int n_mix, int T, float eps, float* mixes); +bool hc_mix_scale(const float* h, int hcdim, int n_mix, int T, float eps, float* mixes); +bool hc_post(const float* y, float* h, const float* post, const float* comb, int dim, int hc, int T); +bool hc_post_v(const float* y, const float* res, const float* post, const float* comb, int dim, int hc, int T, float* out); +bool hc_pre_norm(const float* h, const float* mix, const float* w, int dim, int hc, int T, float eps, float* x); +bool hc_split(const float* mixes, const float* scale, const float* base, int hc, int iters, float hc_eps, int T, float* pre, float* post, float* comb); +bool hc_mix_split(const float* h, const float* fn, int hcdim, int n_mix, int T, float eps, float* mixes, const float* scale, const float* base, int hc, int iters, float hc_eps, float* pre, float* post, float* comb); +bool index_score(const float* q, const float* wq, const void* kc, int heads, int hd, int T, int n, const int* cand, int cand_stride, int blk, const void* rows, float* out, long long out_stride, int kfmt); +bool init_pre(float* pre, int T, void*); +bool iota_rows(int* out, int T, int n, int stride); +bool iq_dn(const unsigned long long* plan_ptr, const unsigned long long* grp_ptr, const int* grp_start, const int* grp_ent, const int* n_groups, const size_t* tr, int n_embd, int n_ff, const int8_t* hq, const float* hd, float* part, int type, int groups); +bool iq_gu(const unsigned long long* plan_ptr, const unsigned long long* grp_ptr, const int* grp_start, const int* grp_ent, const int* n_groups, const size_t* tr, int n_embd, int n_ff, float limit, const int8_t* xq, const float* xd, float* h, int type, int groups); +bool iq_mv(const void* jobs, int nj); +bool markov_add(const float* logits, const uint16_t* E, const uint16_t* H, const int* tok, int V, int R); +bool moe_combine(void* acc, const void* y, const int* pos, const float* w_, int T, int k, int r0, int r1, int dim); +bool moe_sort(const int* idx, int T, int k, const int* vs, int n_vs, int* bounds, int* row_tok, int* pos); +bool moe_win_group_cpu(const void* plan_ptr, void* grp_ptr, void* grp_start, void* grp_ent, void* n_groups, int T, int k); +bool mx_combine(const void* plan_ptr, const void* plan_w, const void* part, const void* cpu_part, size_t cpu_stride, const int* cpu_flag, const void* sh_out, void* out, size_t out_stride, int T, int k, int n_embd); +bool mx_dn(const unsigned long long* plan_ptr, const unsigned long long* grp_ptr, const int* grp_start, const int* grp_ent, const int* n_groups, const size_t* tr, int n_embd, int n_ff, const int8_t* hq, const float* hd, float* part, int groups); +bool mx_gu(const unsigned long long* plan_ptr, const unsigned long long* grp_ptr, const int* grp_start, const int* grp_ent, const int* n_groups, const size_t* tr, int n_embd, int n_ff, float limit, const int8_t* xq, const float* xd, float* h, int groups); +bool mx_mv(const void* jobs, int nj, int max_n, size_t job_bytes, int fmt); +bool q8_quant_rows(const float* x, size_t ldx, int rows, int k, int8_t* q, float* d); +bool ring_to_ext(const void* ring, void* ext, long long p0, int n, int d, size_t row_bytes); +bool route(const float* logits, const float* bias, int n, int k, int T, float route_scale, int* idx, float* w, const float* bias_vl, const void* vl); +bool route_publish(const float* logits, const float* bias, int n, int k, int T, float route_scale, const void* vram, int* idx, float* w, void* plan_ptr, const float* bias_vl, unsigned vl); +bool sample_greedy(const float* logits, int n_tokens, int n_vocab, const int* history, int history_len, int penalty_last_n, float penalty_repeat, float penalty_freq, float penalty_present, const float* logit_bias, int* out); +bool sgemm_tn(const float* A, const float* B, float* C, int m, int nt, int k, int lda, int ldb, int ldc); +bool sparse_attn(const float* q, const uint16_t* win, const void* ckv, const int* sel, int sel_stride, const void* rows, const float* sink, int heads, int d, int T, float scale, float* out, int ckv_fmt); +bool swiglu(const float* g, const float* u, int n, float lim, float* h); +bool topk(const float* vals, long long stride, int T, const void* rows, const int* cand, int cand_stride, int blk, int* out, int out_stride); +bool unpack_sel(const void* ckv, const int* sel, int sel_stride, int gen, long long p0, int t0, int R, int kcap, const int* nsel, int cap, void* out, int T, int nsel_count, size_t row_bytes); + +} // namespace strata::dsv41::metal_host diff --git a/edge0-strata/metal/include/metal_compat/immintrin.h b/edge0-strata/metal/include/metal_compat/immintrin.h new file mode 100644 index 00000000..1f39eddd --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/immintrin.h @@ -0,0 +1,5 @@ +// metal/include/metal_compat/immintrin.h - the x86 intrinsics the host loops use, on arm64: the spin-wait hint and +// the store fence before a host flag is raised (verify.cpp, session.cpp, mtp.cpp, graph.cpp). +#pragma once +static inline void _mm_pause() { __asm__ __volatile__("yield"); } +static inline void _mm_sfence() { __asm__ __volatile__("dmb ishst" ::: "memory"); } diff --git a/edge0-strata/metal/include/metal_compat/linux/mman.h b/edge0-strata/metal/include/metal_compat/linux/mman.h new file mode 100644 index 00000000..5b90ec33 --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/linux/mman.h @@ -0,0 +1,12 @@ +// metal/include/metal_compat/linux/mman.h - the Linux mmap flags src/core/pinned.cu asks for, on macOS. There are no +// hugetlb pages and no transparent huge pages: MAP_HUGETLB is spelled as MAP_FIXED at address 0, which macOS refuses +// (the 4 GiB __PAGEZERO), and MADV_HUGEPAGE as an advice madvise refuses, so pinned.cu takes its 4 KB-page fallback +// and reports it, exactly as on a Linux box without a hugetlb pool. +#pragma once +#include +#ifndef MAP_HUGETLB +#define MAP_HUGETLB MAP_FIXED +#endif +#ifndef MADV_HUGEPAGE +#define MADV_HUGEPAGE (-1) +#endif diff --git a/edge0-strata/metal/include/metal_compat/mma.h b/edge0-strata/metal/include/metal_compat/mma.h new file mode 100644 index 00000000..511071f0 --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/mma.h @@ -0,0 +1,38 @@ +// mma.h - the wmma surface on Metal: the names parse, and using one stops the run. +// +// The DeepSeek token path does not use tensor cores. The prompt path has an MMA variant written for CUDA/HIP warp matrices; +// Metal's equivalent is the simdgroup matrix API, which is a different shape, so that variant is a later piece of the port. +// These aliases exist so the shared sources compile (which is how the rest of the API surface is enumerated), and every +// operation aborts with a message naming this file - an accidental call cannot return plausible garbage. +#pragma once + +#include +#include + +namespace nvcuda { +namespace wmma { + +// CUDA 12 spells these as tags: fragment +struct matrix_a {}; +struct matrix_b {}; +struct accumulator {}; +struct row_major {}; +struct col_major {}; +enum layout_t { mem_row_major, mem_col_major }; + +template struct fragment { + T x[M * N]; +}; + +[[noreturn]] inline void strata_metal_wmma_missing(const char* what) { + std::fprintf(stderr, "strata metal: wmma::%s has no Metal implementation (simdgroup matrices are the equivalent)\n", what); + std::abort(); +} + +template inline void fill_fragment(A&&...) { strata_metal_wmma_missing("fill_fragment"); } +template inline void load_matrix_sync(A&&...) { strata_metal_wmma_missing("load_matrix_sync"); } +template inline void mma_sync(A&&...) { strata_metal_wmma_missing("mma_sync"); } // variadic: the engine passes four fragments of different types +template inline void store_matrix_sync(A&&...) { strata_metal_wmma_missing("store_matrix_sync"); } + +} // namespace wmma +} // namespace nvcuda diff --git a/edge0-strata/metal/include/metal_compat/vector_types.h b/edge0-strata/metal/include/metal_compat/vector_types.h new file mode 100644 index 00000000..2dd9c14d --- /dev/null +++ b/edge0-strata/metal/include/metal_compat/vector_types.h @@ -0,0 +1,25 @@ +// metal/include/metal_compat/vector_types.h - the CUDA vector types the dsv41 sources name, on Metal. +// +// The compat directory had none, so the first half-precision header to use float2 failed. Only the shapes the engine +// touches are here (float2/float4/int2/uint2 and their make_ helpers); the Metal runtime's own float4 is not involved. +#pragma once + +struct float2 { float x, y; }; +struct float4 { float x, y, z, w; }; +struct int2 { int x, y; }; +struct uint2 { unsigned x, y; }; + +static inline float2 make_float2(float x, float y) { return float2{x, y}; } +static inline float4 make_float4(float x, float y, float z, float w) { return float4{x, y, z, w}; } +static inline int2 make_int2(int x, int y) { return int2{x, y}; } +static inline uint2 make_uint2(unsigned x, unsigned y) { return uint2{x, y}; } + +struct int4 { int x, y, z, w; }; +struct uint4 { unsigned x, y, z, w; }; +struct char4 { signed char x, y, z, w; }; +struct uchar4 { unsigned char x, y, z, w; }; +struct short2 { short x, y; }; +struct int3 { int x, y, z; }; + +static inline int4 make_int4(int x, int y, int z, int w) { return int4{x, y, z, w}; } +static inline uint4 make_uint4(unsigned x, unsigned y, unsigned z, unsigned w) { return uint4{x, y, z, w}; } diff --git a/edge0-strata/metal/include/strata/gpu/metal.hpp b/edge0-strata/metal/include/strata/gpu/metal.hpp new file mode 100644 index 00000000..4880bdc8 --- /dev/null +++ b/edge0-strata/metal/include/strata/gpu/metal.hpp @@ -0,0 +1,44 @@ +// metal/include/strata/gpu/metal.hpp - what the Metal kernel library needs beyond include/strata/gpu/runtime.hpp. +// +// The host code never includes this: it sees only runtime.hpp and the strata::kernels::* headers. The Metal +// implementations of those kernel headers (metal/src/kernels/*) turn their `void* stream` into a gpu::Stream and call +// dispatch(): encoded now, or recorded when the stream is being captured (the Metal "graph"). +#pragma once + +#include +#include +#include + +#include "strata/gpu/runtime.hpp" + +namespace MTL { +class Device; +class ComputePipelineState; +} // namespace MTL + +namespace strata::gpu::metal { + +MTL::Device* device(); + +/// Compile MSL source at run time (the Command Line Tools have no offline `metal` compiler) and return the pipeline +/// of `function`. Cached by (source name, function name); nullptr and `err` on failure. +MTL::ComputePipelineState* pipeline(const char* source_name, const char* source, const char* function, std::string& err); + +struct Grid { + uint32_t x = 1, y = 1, z = 1; +}; + +/// One compute dispatch: `args` (<= 4 KiB, copied now, as a CUDA kernel's arguments are at launch/capture) is bound at +/// buffer index 0; the kernel reads its device pointers (gpuAddress values) from it. `threadgroups` x `threads` +/// like a CUDA <<>>. Dispatches on one stream are ordered (each one sees the previous one's +/// writes), as kernels on one CUDA stream are. +Error dispatch(Stream s, MTL::ComputePipelineState* p, const void* args, size_t arg_bytes, Grid threadgroups, Grid threads, + size_t threadgroup_bytes = 0); + +/// Counters for tests and the dispatch-cost budget. +struct Stats { + uint64_t dispatches = 0, command_buffers = 0, blits = 0, host_syncs = 0; +}; +Stats stats(); + +} // namespace strata::gpu::metal diff --git a/edge0-strata/metal/include/strata/gpu/runtime.hpp b/edge0-strata/metal/include/strata/gpu/runtime.hpp new file mode 100644 index 00000000..25ec0c1e --- /dev/null +++ b/edge0-strata/metal/include/strata/gpu/runtime.hpp @@ -0,0 +1,127 @@ +// metal/include/strata/gpu/runtime.hpp - the device runtime of the Metal build (MAC-PLAN Phase 1, compat design). +// +// The shared host files keep calling the CUDA runtime names; on Metal those names come from +// metal/include/metal_compat/cuda_runtime.h (first on the include path, as include/strata/hip_compat is on HIP), which +// forwards every one to the functions declared here. metal/src/gpu/runtime.cpp implements them on metal-cpp: +// - a device pointer is the MTLBuffer's gpuAddress (the host never dereferences it; kernels take it from their +// argument struct); host_alloc returns contents() and host_device_pointer maps it to its gpuAddress; every buffer is +// in one residency set; +// - a Stream is a command queue + the open command buffer / compute encoder + a shared-event timeline; an Event is a +// (shared event, value) pair; +// - begin_capture / end_capture / graph_launch record the stream's dispatches and re-encode them per launch (the CUDA +// graph rule "per-step values live in device buffers, not in arguments" holds here too); +// - caps().unified_memory is true and host_view(device_pointer) is the CPU address of a device allocation. +// The CUDA and HIP builds never include this directory. +#pragma once + +#include +#include + +// The signatures and the meaning of each call are those of the CUDA runtime call of the same purpose. + +namespace strata::gpu { + +enum Error : int { kSuccess = 0, kNotReady = 600, kErrorInvalidValue = 1, kErrorMemoryAllocation = 2, kErrorUnsupported = 801, + kErrorLaunchFailure = 719 }; +struct StreamImpl; +struct EventImpl; +struct GraphImpl; +struct GraphExecImpl; +using Stream = StreamImpl*; // nullptr = the default (legacy) stream, as on CUDA +using Event = EventImpl*; +using Graph = GraphImpl*; +using GraphExec = GraphExecImpl*; +enum CopyKind : int { kHostToHost = 0, kHostToDevice = 1, kDeviceToHost = 2, kDeviceToDevice = 3, kCopyDefault = 4 }; +using HostFn = void (*)(void*); +struct DeviceProp { // the fields the host code reads + char name[256] = {}; + char gcnArchName[256] = {}; + int major = 0, minor = 0, warpSize = 32, multiProcessorCount = 0; + size_t totalGlobalMem = 0; + int integrated = 0; // the DeepSeek engine reads it (a GPU on the PC's RAM: Strix Halo); 0 here: its rules for that part were not tuned for Apple GPUs +}; + +inline constexpr unsigned kHostAllocDefault = 0; +inline constexpr unsigned kHostAllocPortable = 1; +inline constexpr unsigned kHostAllocMapped = 2; +inline constexpr unsigned kHostRegisterPortable = 1; +inline constexpr unsigned kHostRegisterMapped = 2; +inline constexpr unsigned kHostRegisterReadOnly = 8; +inline constexpr unsigned kEventDefault = 0; +inline constexpr unsigned kEventBlockingSync = 1; +inline constexpr unsigned kEventDisableTiming = 2; +inline constexpr unsigned kStreamNonBlocking = 1; + +struct Caps { + bool unified_memory = true; +}; +inline Caps caps() { return Caps{}; } +void* host_view(const void* device_pointer); + +Error get_last_error(); +Error peek_last_error(); +const char* error_string(Error e); + +Error get_device(int* dev); +Error set_device(int dev); +Error device_count(int* n); +Error device_synchronize(); +Error mem_get_info(size_t* free_b, size_t* total_b); +Error get_device_properties(DeviceProp* p, int dev); +Error device_integrated(int* v, int dev); +Error device_sm_count(int* v, int dev); +Error device_clock_khz(int* v, int dev); +Error runtime_version(int* v); +Error device_can_access_peer(int* v, int dev, int peer); +Error device_enable_peer_access(int peer, unsigned flags); +inline constexpr Error kErrorPeerAccessAlreadyEnabled = Error(704); + +Error malloc_device_raw(void** p, size_t bytes); +template inline Error malloc_device(T** p, size_t bytes) { return malloc_device_raw((void**) (void*) p, bytes); } +Error free_device(void* p); +Error host_alloc_raw(void** p, size_t bytes, unsigned flags); +template inline Error host_alloc(T** p, size_t bytes, unsigned flags) { return host_alloc_raw((void**) (void*) p, bytes, flags); } +Error free_host(void* p); +template inline Error malloc_host(T** p, size_t bytes) { return host_alloc_raw((void**) (void*) p, bytes, 0); } +Error host_device_pointer_raw(void** dev, void* host); +template inline Error host_device_pointer(T** dev, void* host, unsigned = 0) { + return host_device_pointer_raw((void**) (void*) dev, host); +} +Error host_register(void* p, size_t bytes, unsigned flags); +Error host_unregister(void* p); +Error memset_device(void* p, int value, size_t bytes); +Error memset_async(void* p, int value, size_t bytes, Stream s = nullptr); + +Error memcpy_sync(void* dst, const void* src, size_t bytes, CopyKind k); +Error memcpy_async(void* dst, const void* src, size_t bytes, CopyKind k, Stream s = nullptr); +Error memcpy_2d_async(void* dst, size_t dpitch, const void* src, size_t spitch, size_t width, size_t height, CopyKind k, + Stream s = nullptr); +Error memcpy_peer_async(void* dst, int dst_dev, const void* src, int src_dev, size_t bytes, Stream s = nullptr); + +Error stream_create(Stream* s, unsigned flags); +inline Error stream_create(Stream* s) { return stream_create(s, 0); } +Error stream_destroy(Stream s); +Error stream_synchronize(Stream s); +Error stream_query(Stream s); +Error stream_wait_event(Stream s, Event e, unsigned flags = 0); +Error launch_host_func(Stream s, HostFn fn, void* arg); + +Error event_create(Event* e, unsigned flags); +inline Error event_create(Event* e) { return event_create(e, 0); } +Error event_destroy(Event e); +Error event_record(Event e, Stream s = nullptr); +Error event_query(Event e); +Error event_synchronize(Event e); +Error event_elapsed_ms(float* ms, Event a, Event b); + +Error begin_capture(Stream s); +Error end_capture(Stream s, Graph* g); +Error graph_node_count(Graph g, size_t* n); +Error graph_instantiate(GraphExec* x, Graph g); +inline Error graph_instantiate(GraphExec* x, Graph g, unsigned long long) { return graph_instantiate(x, g); } +Error graph_launch(GraphExec x, Stream s); +Error graph_upload(GraphExec x, Stream s); +Error graph_destroy(Graph g); +Error graph_exec_destroy(GraphExec x); + +} // namespace strata::gpu diff --git a/edge0-strata/metal/kernels/alu_bench.metal b/edge0-strata/metal/kernels/alu_bench.metal new file mode 100644 index 00000000..e34890a0 --- /dev/null +++ b/edge0-strata/metal/kernels/alu_bench.metal @@ -0,0 +1,181 @@ +// metal/kernels/alu_bench.metal - instruction-rate probes (alu_bench.cpp): each thread runs ITER iterations of 8 +// independent chains of one operation; the result is written so nothing is eliminated. +#define ITER 1024 + +kernel void a_fma32(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(a0, k, 0.5f); a1 = fma(a1, k, 0.5f); a2 = fma(a2, k, 0.5f); a3 = fma(a3, k, 0.5f); + a4 = fma(a4, k, 0.5f); a5 = fma(a5, k, 0.5f); a6 = fma(a6, k, 0.5f); a7 = fma(a7, k, 0.5f); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +kernel void a_fma16(device float* out [[buffer(0)]], constant float& kf [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + const half k = half(kf); + half a0 = gid & 7, a1 = 1, a2 = 2, a3 = 3, a4 = 4, a5 = 5, a6 = 6, a7 = 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(a0, k, half(0.5)); a1 = fma(a1, k, half(0.5)); a2 = fma(a2, k, half(0.5)); a3 = fma(a3, k, half(0.5)); + a4 = fma(a4, k, half(0.5)); a5 = fma(a5, k, half(0.5)); a6 = fma(a6, k, half(0.5)); a7 = fma(a7, k, half(0.5)); + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +// float accumulator, half operands (does the conversion come free?) +kernel void a_fmamix(device float* out [[buffer(0)]], constant float& kf [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half h0 = half(gid & 7), h1 = h0 + half(1); + float a0 = gid, a1 = 1, a2 = 2, a3 = 3, a4 = 4, a5 = 5, a6 = 6, a7 = 7; + for (int i = 0; i < ITER; ++i) { + a0 = fma(float(h0), kf, a0); a1 = fma(float(h1), kf, a1); a2 = fma(float(h0), kf, a2); a3 = fma(float(h1), kf, a3); + a4 = fma(float(h0), kf, a4); a5 = fma(float(h1), kf, a5); a6 = fma(float(h0), kf, a6); a7 = fma(float(h1), kf, a7); + h0 = h0 * half(0.999); h1 = h1 * half(0.999); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +kernel void a_iand(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = (a0 ^ k) + 1; a1 = (a1 ^ k) + 1; a2 = (a2 ^ k) + 1; a3 = (a3 ^ k) + 1; + a4 = (a4 ^ k) + 1; a5 = (a5 ^ k) + 1; a6 = (a6 ^ k) + 1; a7 = (a7 ^ k) + 1; + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +kernel void a_shift(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint a0 = gid, a1 = gid + 1, a2 = gid + 2, a3 = gid + 3, a4 = gid + 4, a5 = gid + 5, a6 = gid + 6, a7 = gid + 7; + for (int i = 0; i < ITER; ++i) { + a0 = (a0 >> k) | k; a1 = (a1 >> k) | k; a2 = (a2 >> k) | k; a3 = (a3 >> k) | k; + a4 = (a4 >> k) | k; a5 = (a5 >> k) | k; a6 = (a6 >> k) | k; a7 = (a7 >> k) | k; + } + out[gid] = float(a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7); +} + +// int -> float conversion + fma +kernel void a_cvt(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid; + float a0 = 0, a1 = 0, a2 = 0, a3 = 0, a4 = 0, a5 = 0, a6 = 0, a7 = 0; + for (int i = 0; i < ITER; ++i) { + u += k; + a0 += float(u & 3); a1 += float(u & 12); a2 += float(u & 48); a3 += float(u & 192); + a4 += float(u & 768); a5 += float(u & 3072); a6 += float(u & 12288); a7 += float(u & 49152); + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +// threadgroup memory loads: random-ish 8-byte reads from a 4 KB table +kernel void a_tgload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint2 tab[512]; + for (int i = tid; i < 512; i += 256) tab[i] = uint2(i, i * 3); + threadgroup_barrier(mem_flags::mem_threadgroup); + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} + +// select (sign application) +kernel void a_sel(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid; + float a0 = 1, a1 = 2, a2 = 3, a3 = 4, a4 = 5, a5 = 6, a6 = 7, a7 = 8; + for (int i = 0; i < ITER; ++i) { + u ^= k; + a0 = (u & 1) ? -a0 : a0; a1 = (u & 2) ? -a1 : a1; a2 = (u & 4) ? -a2 : a2; a3 = (u & 8) ? -a3 : a3; + a4 = (u & 16) ? -a4 : a4; a5 = (u & 32) ? -a5 : a5; a6 = (u & 64) ? -a6 : a6; a7 = (u & 128) ? -a7 : a7; + } + out[gid] = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7; +} + +// dot-shaped: acc_i += q_i * x_i with loop-carried accumulators only (q, x vary slowly) +kernel void a_dot32(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float q[8], x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = float(gid + j); x[j] = k * j; a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(q[j], x[j], a[j]); + q[0] += 1.0f; + } + float s = 0; + for (int j = 0; j < 8; ++j) s += a[j]; + out[gid] = s; +} +kernel void a_dot16(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half q[8], x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = half(gid + j); x[j] = half(k * j); a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(q[j], x[j], a[j]); + q[0] += half(1.0f); + } + float s = 0; + for (int j = 0; j < 8; ++j) s += float(a[j]); + out[gid] = s; +} +kernel void a_dotmix(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + half q[8]; + float x[8], a[8]; + for (int j = 0; j < 8; ++j) { q[j] = half(gid + j); x[j] = k * j; a[j] = 0; } + for (int i = 0; i < ITER; ++i) { + for (int j = 0; j < 8; ++j) a[j] = fma(float(q[j]), x[j], a[j]); + q[0] += half(1.0f); + } + float s = 0; + for (int j = 0; j < 8; ++j) s += a[j]; + out[gid] = s; +} +kernel void a_fma32x16(device float* out [[buffer(0)]], constant float& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + float a[16]; + for (int j = 0; j < 16; ++j) a[j] = gid + j; + for (int i = 0; i < ITER / 2; ++i) + for (int j = 0; j < 16; ++j) a[j] = fma(a[j], k, 0.5f); + float s = 0; + for (int j = 0; j < 16; ++j) s += a[j]; + out[gid] = s; +} + +// 8-byte gathers from a 4 KB table in device memory (L1-resident) and in the constant address space +kernel void a_devload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], device const uint2* tab [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} +kernel void a_cstload(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], constant uint2* tab [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint u = gid * 2654435761u; + uint2 s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 511]; s1 += tab[(u >> 7) & 511]; s2 += tab[(u >> 11) & 511]; s3 += tab[(u >> 15) & 511]; + u += k; + } + out[gid] = float(s0.x + s1.x + s2.x + s3.x + s0.y + s1.y + s2.y + s3.y); +} +// 4-byte gathers from a 1 KB threadgroup table (IQ3_XXS's grid) +kernel void a_tgload4(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint tab[256]; + for (int i = tid; i < 256; i += 256) tab[i] = i * 3; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint u = gid * 2654435761u; + uint s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += tab[(u >> 3) & 255]; s1 += tab[(u >> 7) & 255]; s2 += tab[(u >> 11) & 255]; s3 += tab[(u >> 15) & 255]; + u += k; + } + out[gid] = float(s0 + s1 + s2 + s3); +} +// popcount parity (computed ksigns) +kernel void a_popc(device float* out [[buffer(0)]], constant uint& k [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + uint u = gid, s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int i = 0; i < ITER; ++i) { + s0 += popcount(u & 127); s1 += popcount((u >> 7) & 127); s2 += popcount((u >> 14) & 127); s3 += popcount((u >> 21) & 127); + u += k; + } + out[gid] = float(s0 + s1 + s2 + s3); +} diff --git a/edge0-strata/metal/kernels/dense_mm.metal b/edge0-strata/metal/kernels/dense_mm.metal new file mode 100644 index 00000000..3e149bfb --- /dev/null +++ b/edge0-strata/metal/kernels/dense_mm.metal @@ -0,0 +1,77 @@ +// metal/kernels/dense_mm.metal - dense projections for a PROMPT chunk (after quant/units/mma.metal): +// act_f16 X[t][k] -> half (the matrix units' operand) +// dq_f16 a quantized matrix -> half rows (ref_quant's values rounded to fp16), once per chunk +// gemm_f16 Y[t][n] = Xh[t] . Wh[n] (fp32) by MetalPerformancePrimitives matmul2d, TM x TN tiles, K dynamic; +// rows of Y past n_out / t past n_tok are not written (the tensors carry the extents) +#include +using namespace mpp::tensor_ops; + +kernel void act_f16(device const float* x [[buffer(0)]], device half* xh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (gid < n) xh[gid] = half(x[gid]); +} + +struct DqArgs { + uint n_in; // values per row + uint n_out; // rows + uint row_bytes; + uint perm; // 0: row r -> r; p > 0: row r = c * (n_out / p) + d -> d * p + c (the HC up matrix with the p + // streams of an output interleaved, for gemm_hc_up's epilogue) +}; + +// one thread per unit: grid (units per row / 64, n_out), 64 threads +template +kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer(1)]], constant DqArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 64); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int u = int(tg.x) * 64 + tid, r = int(tg.y); + if (u * UNIT >= int(a.n_in)) return; + // the unit's chunk origin: a 128-wide chunk holds whole units of every format + const int k0 = (u * UNIT) / 128 * 128; + float w[UNIT]; + int p0, p1; + DW::run(W + ulong(r) * a.row_bytes, u, k0, w, p0, p1, tab); + const int q = a.perm ? int(a.n_out / a.perm) : 1; + const int ro = a.perm ? (r % q) * int(a.perm) + r / q : r; + device half* o = Wh + ulong(ro) * a.n_in + k0; + for (int j = 0; j < H; ++j) { + o[p0 + j] = half(w[j]); + o[p1 + j] = half(w[H + j]); + } +} + +struct GemmArgs { int M, N, K; int pad; }; // Y is M rows of N floats + +template +kernel void gemm_f16(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device float* Y [[buffer(2)]], + constant GemmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + tensor, tensor_inline> tY(Y, dextents(a.N, a.M)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mB = tB.slice(0, int(tg.x) * TN); + auto mY = tY.slice(int(tg.x) * TN, int(tg.y) * TM); + op.run(mA, mB, mY); +} + +#define DQ(F) template [[host_name("dq_f16_" #F)]] kernel decltype(dq_f16) dq_f16; +#define DQ_(F) DQ(F) +DQ_(T_F32) DQ_(T_BF16) DQ_(T_Q8_0) DQ_(T_Q4_0) DQ_(T_Q5_0) DQ_(T_IQ4_NL) DQ_(T_IQ4_XS) DQ_(T_Q2_0) DQ_(T_IQ3_S) DQ_(T_Q3_K) +DQ_(T_Q4_K) DQ_(T_Q5_K) DQ_(T_Q6_K) +template [[host_name("gemm_f16_64_64_4")]] kernel decltype(gemm_f16<64, 64, 4>) gemm_f16<64, 64, 4>; +template [[host_name("gemm_f16_128_64_4")]] kernel decltype(gemm_f16<128, 64, 4>) gemm_f16<128, 64, 4>; +template [[host_name("gemm_f16_64_128_4")]] kernel decltype(gemm_f16<64, 128, 4>) gemm_f16<64, 128, 4>; +template [[host_name("gemm_f16_64_32_4")]] kernel decltype(gemm_f16<64, 32, 4>) gemm_f16<64, 32, 4>; +template [[host_name("gemm_f16_128_128_4")]] kernel decltype(gemm_f16<128, 128, 4>) gemm_f16<128, 128, 4>; +template [[host_name("gemm_f16_128_128_8")]] kernel decltype(gemm_f16<128, 128, 8>) gemm_f16<128, 128, 8>; +template [[host_name("gemm_f16_128_64_8")]] kernel decltype(gemm_f16<128, 64, 8>) gemm_f16<128, 64, 8>; +template [[host_name("gemm_f16_64_64_8")]] kernel decltype(gemm_f16<64, 64, 8>) gemm_f16<64, 64, 8>; +template [[host_name("gemm_f16_256_64_8")]] kernel decltype(gemm_f16<256, 64, 8>) gemm_f16<256, 64, 8>; +template [[host_name("gemm_f16_64_256_8")]] kernel decltype(gemm_f16<64, 256, 8>) gemm_f16<64, 256, 8>; +template [[host_name("gemm_f16_128_32_4")]] kernel decltype(gemm_f16<128, 32, 4>) gemm_f16<128, 32, 4>; +template [[host_name("gemm_f16_32_128_4")]] kernel decltype(gemm_f16<32, 128, 4>) gemm_f16<32, 128, 4>; diff --git a/edge0-strata/metal/kernels/gdn.metal b/edge0-strata/metal/kernels/gdn.metal new file mode 100644 index 00000000..20eb4463 --- /dev/null +++ b/edge0-strata/metal/kernels/gdn.metal @@ -0,0 +1,185 @@ +// metal/kernels/gdn.metal - one GDN (gated delta-net) layer's recurrence for a window of tokens, fused end to end: +// +// conv (4 taps, history [channels][3] oldest first) -> SiLU, on q (16 heads), k (16), v (48 heads) of 128 +// q, k: L2 norm as native_gdn_l2_norm (rsqrt(sum/128 + eps/128), then x 1/sqrt(128), two roundings) +// gate = exp(softplus(alpha + dt) * a) (softplus = v > 20 ? v : log1p(exp(v))), beta = sigmoid(beta) +// kv_j = sum_i S[i][j] k[i]; delta_j = (v_j - gate * kv_j) * beta; S = gate * S + k[i] delta_j +// o_j = (sum_i S[i][j] q[i]) / sqrt(128); y_j = o_j * rsqrt(mean(o^2) + eps) * gamma_j * sigmoid(z_j) +// +// State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k. +// One threadgroup per value head, 512 threads = 16 simdgroups: simdgroup s holds columns j = 32 (s % 4) + lane and +// rows i = 32 (s / 4) .. + 32 of the head's 128 x 128 state in registers. GDN_NTOK tokens run in order; with +// `commit` the state and the conv history are written back after the last one (a verify window runs without commit, +// the fold of m accepted tokens is the same kernel with GDN_NTOK = m and commit). The conv history is read from +// `hist` and the committed one written to `hist_out` (ping-pong: the q/k channels are shared by three threadgroups). + +constant int GDN_NTOK [[function_constant(2)]]; +// GDN_ROWS (function constant 31, with GDN_NTOK = 1): the grid is (h_v, rows) and row y is a DIFFERENT sequence (a batch slot) with its own +// state and conv history; rinfo as in qsa.metal (header [5] state stride, [6] history stride per slot and parity; row y at [8 + 4 y] = +// {slot, cell, history parity, 0}). hist / hist_out are then one buffer (hist), the parities side by side per slot. The inputs and the +// output are row y of the window's arrays. Off: exactly the plain kernel. +constant bool GDN_ROWS_C [[function_constant(31)]]; +constant bool gdn_rows = is_function_constant_defined(GDN_ROWS_C) && GDN_ROWS_C; + +struct GdnArgs { + uint h_k; // 16 + uint h_v; // 48 + uint channels; // 10240 = (2 h_k + h_v) * 128 + uint commit; // 0: verify (no write); 1: commit after the last token; 0x100 | k (1 <= k <= GDN_NTOK): commit + // the state and the conv history after the first k tokens, in place, while every token's y is + // still computed (a window whose first k rows are kept) - bitwise the commit of a k-token step + float eps; // 1e-6 + uint plane; // 0: y only; else y is the raw plane of prepared activations (xprep layout, n = h_v * 128, + // `plane` floats between planes) and the other planes and the per-16 sums are written too +}; + +// The kernel is launched with IB * 128 threads (512 for gdn_step, 1024 for gdn_step8). Without the attribute the compiler +// is free to use as many registers as the 128-row state needs and a pipeline can end up with +// maxTotalThreadsPerThreadgroup below the launch size: an M2 Max gave 448 for the GDN_NTOK >= 2 specializations, and +// Metal then skips the dispatch without an error (#2323). The attribute makes the compiler fit the launch size. +template +[[max_total_threads_per_threadgroup(IB * 128)]] +kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[buffer(1)]], + device const float* qkv [[buffer(2)]], device const float* alpha [[buffer(3)]], + device const float* beta [[buffer(4)]], device const float* z [[buffer(5)]], + device const float* conv_w [[buffer(6)]], device const float* dt [[buffer(7)]], + device const float* ssm_a [[buffer(8)]], device const float* gamma [[buffer(9)]], + device float* y [[buffer(10)]], constant GdnArgs& a [[buffer(11)]], + device float* hist_out [[buffer(12)]], + device const uint* rinfo [[buffer(13), function_constant(gdn_rows)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128; + threadgroup float qs[S], ks[S], vs[S]; + constexpr int RPT = S / IB; // state rows per thread + threadgroup float red[IB][S]; + threadgroup float part[32]; + threadgroup float gb[2]; + threadgroup float ys[S]; + const int h = int(tg.x), hk = h % int(a.h_k), hv = int(a.h_v); + const int cb = sg & 3, ib = sg >> 2; + const int j = 32 * cb + lane, i0 = RPT * ib; + const float scale = 1.0f / sqrt(float(S)); + const int row0 = gdn_rows ? int(tg.y) : 0; + if (gdn_rows) { + const ulong sl = rinfo[8 + 4 * row0], par = rinfo[8 + 4 * row0 + 2]; + state += sl * rinfo[5]; + hist_out = hist + (sl * 2 + (1 - par)) * rinfo[6]; + hist += (sl * 2 + par) * rinfo[6]; + } + + // the channel this thread convolves: 0..127 q of head hk, 128..255 k of head hk, 256..383 v of head h + int ch = -1; + if (tid < 128) ch = hk * S + tid; + else if (tid < 256) ch = int(a.h_k) * S + hk * S + (tid - 128); + else if (tid < 384) ch = 2 * int(a.h_k) * S + h * S + (tid - 256); + float hst[3] = {0.0f, 0.0f, 0.0f}, cw[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + if (ch >= 0) { + for (int k = 0; k < 3; ++k) hst[k] = hist[ch * 3 + k]; + for (int k = 0; k < 4; ++k) cw[k] = conv_w[ch * 4 + k]; + } + float st[RPT]; + for (int r = 0; r < RPT; ++r) st[r] = state[((i0 + r) * hv + h) * S + j]; + + const int commit_at = a.commit == 1u ? GDN_NTOK : (a.commit >> 8) ? int(a.commit & 0xFFu) : 0; + for (int t = 0; t < GDN_NTOK; ++t) { + // conv + SiLU + const int tr = t + row0; + float cv = 0.0f; + if (ch >= 0) { + const float in = qkv[tr * int(a.channels) + ch]; + const float vals[4] = {hst[0], hst[1], hst[2], in}; + float s = 0.0f; + for (int k = 0; k < 4; ++k) s = fma(vals[k], cw[k], s); + s = s + 0.0f; + cv = s / (1.0f + exp(-s)); + hst[0] = hst[1]; + hst[1] = hst[2]; + hst[2] = in; + } + // L2 norms of q (threads 0..127 = simdgroups 0..3) and k (128..255 = 4..7) + const float ssq = simd_sum(cv * cv); + if (lane == 0) part[sg] = ssq; + if (tid == 0) { + const float v = alpha[tr * hv + h] + dt[h]; + // log1p(e^v) (Kahan: log(u) * e / (u - 1), u = 1 + e; exact where 1 + e rounds to 1) + const float e = exp(v), u = 1.0f + e; + const float sp = v > 20.0f ? v : (u == 1.0f ? e : log(u) * e / (u - 1.0f)); + gb[0] = exp(sp * ssm_a[h]); + gb[1] = 1.0f / (1.0f + exp(-beta[tr * hv + h])); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + const int base = tid < 128 ? 0 : 4; + const float tot = (part[base] + part[base + 1]) + (part[base + 2] + part[base + 3]); + const float rs = rsqrt(tot / float(S) + a.eps / float(S)); + const float nv = (rs * cv) * scale; + if (tid < 128) qs[tid] = nv; + else ks[tid - 128] = nv; + } else if (tid < 384) { + vs[tid - 256] = cv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float g = gb[0], be = gb[1]; + // kv_j + float kv = 0.0f; + for (int r = 0; r < RPT; ++r) kv = fma(st[r], ks[i0 + r], kv); + red[ib][j] = kv; + threadgroup_barrier(mem_flags::mem_threadgroup); + kv = 0.0f; + for (int k = 0; k < IB; ++k) kv += red[k][j]; + const float delta = (vs[j] - g * kv) * be; + float o = 0.0f; + for (int r = 0; r < RPT; ++r) { + st[r] = fma(g, st[r], ks[i0 + r] * delta); + o = fma(st[r], qs[i0 + r], o); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + red[ib][j] = o; + threadgroup_barrier(mem_flags::mem_threadgroup); + // the readout and the output norm: column j's o (every i-block thread of the column holds the same value) + o = 0.0f; + for (int k = 0; k < IB; ++k) o += red[k][j]; + o = o * scale; + const float osq = ib == 0 ? simd_sum(o * o) : 0.0f; + if (ib == 0 && lane == 0) part[cb] = osq; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ib == 0) { + const float tot = (part[0] + part[1]) + (part[2] + part[3]); + const float rs = rsqrt(tot / float(S) + a.eps); + const float zz = z[tr * hv * S + h * S + j]; + const float yv = ((rs * o) * gamma[j]) * (1.0f / (1.0f + exp(-zz))); + y[tr * hv * S + h * S + j] = yv; + if (a.plane) { + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; + const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + device float* raw = y + tr * hv * S + h * S + j; + const int jj = j & 15; + raw[a.plane] = yv * q2s[jj & 7]; + raw[2 * a.plane] = yv * p4s[jj & 3]; + raw[3 * a.plane] = (jj & 1) ? yv * 0.00390625f : yv; + ys[j] = yv; + } + } + if (t + 1 == commit_at) { + // the state and the conv history after token t (registers: st after this token's update, hst after its + // shift); the state is read only before the loop, so in place is safe. q/k history: the threadgroups + // h < h_k own their q/k head (h == hk); v history: every head its own - to hist_out, NOT in place: a q/k + // channel is read by the three threadgroups h % h_k at their start, and a threadgroup may start after its + // owner has finished (seen on the 10-core M4); the caller swaps the two + for (int r = 0; r < RPT; ++r) state[((i0 + r) * hv + h) * S + j] = st[r]; + if (ch >= 0 && (tid >= 256 || h < int(a.h_k))) + for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (a.plane && ib == 0 && (j & 15) == 0) { + float s0 = 0.0f, s1 = 0.0f; + for (int k = 0; k < 16; k += 2) { s0 += ys[j + k]; s1 += ys[j + k + 1]; } + y[4 * a.plane + tr * (hv * S / 16) + (h * S + j) / 16] = s0 + s1; + } + } +} + +template [[host_name("gdn_step")]] kernel decltype(gdn_step_t<4>) gdn_step_t<4>; +template [[host_name("gdn_step8")]] kernel decltype(gdn_step_t<8>) gdn_step_t<8>; diff --git a/edge0-strata/metal/kernels/gdn_prompt.metal b/edge0-strata/metal/kernels/gdn_prompt.metal new file mode 100644 index 00000000..43b42aa4 --- /dev/null +++ b/edge0-strata/metal/kernels/gdn_prompt.metal @@ -0,0 +1,789 @@ +// metal/kernels/gdn_prompt.metal - a GDN layer's recurrence for a PROMPT chunk of T tokens, in three passes: the work +// that does not touch the state runs parallel over the tokens; only the state update is a scan, and the columns of a +// head's state are independent in it (column j needs k, q, v_j, the gate and beta), so a lane owns one column. +// +// gdn_pre per (token, head): conv (4 taps; history oldest first) -> SiLU; q, k L2 norm * 1/sqrt(128) (the decode +// step's arithmetic); kq = k . q per key head; gate = exp(softplus(alpha + dt) * a), beta = sigmoid(beta); +// the chunk's last 3 conv inputs -> hist_out +// gdn_scan per (value head, 32 columns x CB simdgroups): the state column in registers (128 floats a lane); +// per token a = S^T k, b = S^T q (old state), delta = (v - g a) beta, o = (g b + delta kq) / sqrt(128) +// (= S_new^T q / sqrt(128)), S = g S + k delta^T; q and k staged per TB tokens in threadgroup memory +// gdn_post per (token, head): y = o * rsqrt(mean(o^2) + eps) * gamma * sigmoid(z) -> fp16 (the out GEMM's operand) +// and, with `f32`, fp32 +// +// State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k. +// The kernels are new entries (nothing of gdn.metal changes). + +struct GdnPArgs { + uint n_tok; + uint h_k; // 16 + uint h_v; // 48 + uint channels; // 10240 + float eps; // 1e-6 + uint f32; // gdn_post: also write fp32 y +}; + +// function constant 7 (optional): the row stride in floats of ONE fused output [qkv | z | alpha | beta] of the four input +// projections (16480). gdn_pre then reads qkv / alpha / beta and gdn_post reads z at that stride (the buffers bound at offsets); +// undefined (every other caller): the separate buffers of channels / h_v / h_v * 128 floats a row. +constant int GDN_LD [[function_constant(7)]]; +constant bool gdn_ld_on = is_function_constant_defined(GDN_LD); + +static inline float gp_conv(device const float* qkv, device const float* hist, device const float* cw, int ch, int t, + int channels) { + float s = 0.0f; + for (int k = 0; k < 4; ++k) { + const int src = t - 3 + k; + const float x = src >= 0 ? qkv[src * channels + ch] : hist[ch * 3 + (src + 3)]; + s = fma(x, cw[ch * 4 + k], s); + } + s = s + 0.0f; + return s / (1.0f + exp(-s)); +} + +// grid (h_k + h_v / 2, T), 256 threads. x < h_k: q (threads 0..127) and k (128..255) of key head x; else value heads +// 2 (x - h_k) and + 1. +kernel void gdn_pre(device const float* qkv [[buffer(0)]], device const float* hist [[buffer(1)]], + device const float* conv_w [[buffer(2)]], device const float* alpha [[buffer(3)]], + device const float* beta [[buffer(4)]], device const float* dt [[buffer(5)]], + device const float* ssm_a [[buffer(6)]], device float* qn [[buffer(7)]], device float* kn [[buffer(8)]], + device float* vc [[buffer(9)]], device float2* gb [[buffer(10)]], device float* kq [[buffer(11)]], + device float* hist_out [[buffer(12)]], constant GdnPArgs& a [[buffer(13)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128; + threadgroup float part[8]; + threadgroup float prod[S]; + const int x = int(tg.x), t = int(tg.y), hk_n = int(a.h_k), hv = int(a.h_v), C = int(a.channels), T = int(a.n_tok); + int ch; + if (x < hk_n) ch = tid < 128 ? x * S + tid : hk_n * S + x * S + (tid - 128); + else ch = 2 * hk_n * S + (2 * (x - hk_n) + (tid >> 7)) * S + (tid & 127); + const int QL = gdn_ld_on ? GDN_LD : C; + const float cv = gp_conv(qkv, hist, conv_w, ch, t, QL); + if (t == T - 1) + for (int k = 0; k < 3; ++k) { + const int src = T - 3 + k; + hist_out[ch * 3 + k] = src >= 0 ? qkv[src * QL + ch] : hist[ch * 3 + (src + 3)]; + } + if (x < hk_n) { + const float ssq = simd_sum(cv * cv); + if (lane == 0) part[sg] = ssq; + threadgroup_barrier(mem_flags::mem_threadgroup); + const int base = tid < 128 ? 0 : 4; + const float tot = (part[base] + part[base + 1]) + (part[base + 2] + part[base + 3]); + const float rs = rsqrt(tot / float(S) + a.eps / float(S)); + const float nv = (rs * cv) * (1.0f / sqrt(float(S))); + if (tid < 128) { + qn[(t * hk_n + x) * S + tid] = nv; + prod[tid] = nv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid >= 128) { + kn[(t * hk_n + x) * S + tid - 128] = nv; + const float p = simd_sum(nv * prod[tid - 128]); + if (lane == 0) part[sg] = p; // simdgroups 4..7 (q's slots 0..3 are read already) + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) kq[t * hk_n + x] = (part[4] + part[5]) + (part[6] + part[7]); + } else { + const int h = 2 * (x - hk_n) + (tid >> 7); + vc[(t * hv + h) * S + (tid & 127)] = cv; + if ((tid & 127) == 0) { + const int AL = gdn_ld_on ? GDN_LD : hv; + const float v = alpha[t * AL + h] + dt[h]; + const float e = exp(v), u = 1.0f + e; + const float sp = v > 20.0f ? v : (u == 1.0f ? e : log(u) * e / (u - 1.0f)); + gb[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * AL + h]))); + } + } +} + +// grid (h_v * 4 / CB), CB * 32 threads; TB tokens of q and k staged at a time +template +kernel void gdn_scan_t(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], + device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, NB = 4 / CB; + threadgroup float4 qs[TB][32], ks[TB][32]; + const int h = int(tg.x) / NB, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int j = ((int(tg.x) % NB) * CB + sg) * 32 + lane; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + float st[S]; +#pragma clang loop unroll(full) + for (int i = 0; i < S; ++i) st[i] = state[(i * hv + h) * S + j]; + for (int t0 = 0; t0 < T; t0 += TB) { + const int nb = min(TB, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += CB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt; + const float2 g_b = gb[t * hv + h]; + const float kqv = kq[t * hkn + hk]; + const float v = vc[(t * hv + h) * S + j]; + float ka = 0.0f, kb = 0.0f, qa = 0.0f, qb = 0.0f; +#pragma clang loop unroll(full) + for (int w = 0; w < 32; w += 2) { + const float4 k0 = ks[tt][w], k1 = ks[tt][w + 1], q0 = qs[tt][w], q1 = qs[tt][w + 1]; + ka = fma(st[4 * w], k0.x, ka); ka = fma(st[4 * w + 1], k0.y, ka); + ka = fma(st[4 * w + 2], k0.z, ka); ka = fma(st[4 * w + 3], k0.w, ka); + kb = fma(st[4 * w + 4], k1.x, kb); kb = fma(st[4 * w + 5], k1.y, kb); + kb = fma(st[4 * w + 6], k1.z, kb); kb = fma(st[4 * w + 7], k1.w, kb); + qa = fma(st[4 * w], q0.x, qa); qa = fma(st[4 * w + 1], q0.y, qa); + qa = fma(st[4 * w + 2], q0.z, qa); qa = fma(st[4 * w + 3], q0.w, qa); + qb = fma(st[4 * w + 4], q1.x, qb); qb = fma(st[4 * w + 5], q1.y, qb); + qb = fma(st[4 * w + 6], q1.z, qb); qb = fma(st[4 * w + 7], q1.w, qb); + } + const float g = g_b.x, kv = ka + kb, qv = qa + qb; + const float delta = (v - g * kv) * g_b.y; + o[(t * hv + h) * S + j] = fma(g, qv, delta * kqv) * scale; +#pragma clang loop unroll(full) + for (int w = 0; w < 32; ++w) { + const float4 k4 = ks[tt][w]; + st[4 * w] = fma(g, st[4 * w], k4.x * delta); + st[4 * w + 1] = fma(g, st[4 * w + 1], k4.y * delta); + st[4 * w + 2] = fma(g, st[4 * w + 2], k4.z * delta); + st[4 * w + 3] = fma(g, st[4 * w + 3], k4.w * delta); + } + } + } +#pragma clang loop unroll(full) + for (int i = 0; i < S; ++i) state[(i * hv + h) * S + j] = st[i]; +} + +// a probe for the bench (function constant 6): 1 = no per-token exchange (wrong results), 2 = no state update +constant int GDN_PROBE [[function_constant(6)]]; +constant int gdn_probe = is_function_constant_defined(GDN_PROBE) ? GDN_PROBE : 0; + +// the rows of a column split over IB simdgroups (128 / IB rows a lane, fewer registers): grid (h_v * 4), IB * 32 +// threads; lane = column (32 columns a threadgroup); per token the IB partial dots are exchanged through threadgroup +// memory (double-buffered: one barrier a token) +template +kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], + device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, R = S / IB, R4 = R / 4; + threadgroup float4 qs[TB][32], ks[TB][32]; + threadgroup float2 red[2][IB][32]; + threadgroup float vs[TB][32]; + threadgroup float4 gk[TB]; // gate, beta, kq + const int h = int(tg.x) / 4, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * R; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + float st[R]; +#pragma clang loop unroll(full) + for (int i = 0; i < R; ++i) st[i] = state[((i0 + i) * hv + h) * S + j]; + for (int t0 = 0; t0 < T; t0 += TB) { + const int nb = min(TB, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += IB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; + vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; + if (w == 0) { + const float2 g_b = gb[(t0 + tt) * hv + h]; + gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt, slot = t & 1; + float ka = 0.0f, qa = 0.0f; +#pragma clang loop unroll(full) + for (int w = 0; w < R4; ++w) { + const float4 k4 = ks[tt][i0 / 4 + w], q4 = qs[tt][i0 / 4 + w]; + ka = fma(st[4 * w], k4.x, ka); ka = fma(st[4 * w + 1], k4.y, ka); + ka = fma(st[4 * w + 2], k4.z, ka); ka = fma(st[4 * w + 3], k4.w, ka); + qa = fma(st[4 * w], q4.x, qa); qa = fma(st[4 * w + 1], q4.y, qa); + qa = fma(st[4 * w + 2], q4.z, qa); qa = fma(st[4 * w + 3], q4.w, qa); + } + const float4 g_b = gk[tt]; + const float kqv = g_b.z; + const float v = vs[tt][lane]; + float2 s = float2(ka, qa); + if (gdn_probe != 1) { + red[slot][sg][lane] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + s = red[slot][0][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][b][lane]; + } + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (sg == 0) o[(t * hv + h) * S + j] = fma(g, s.y, delta * kqv) * scale; + if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int w = 0; w < R4; ++w) { + const float4 k4 = ks[tt][i0 / 4 + w]; + st[4 * w] = fma(g, st[4 * w], k4.x * delta); + st[4 * w + 1] = fma(g, st[4 * w + 1], k4.y * delta); + st[4 * w + 2] = fma(g, st[4 * w + 2], k4.z * delta); + st[4 * w + 3] = fma(g, st[4 * w + 3], k4.w * delta); + } + } + } +#pragma clang loop unroll(full) + for (int i = 0; i < R; ++i) state[((i0 + i) * hv + h) * S + j] = st[i]; +} + +// gdn_scan_r<8, 16> with the 16 state rows of a lane in four float4 variables (no array): a register-allocation check +kernel void gdn_scan_v(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], + device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, IB = 8, TB = 16; + threadgroup float4 qs[TB][32], ks[TB][32]; + threadgroup float2 red[2][IB][32]; + threadgroup float vs[TB][32]; + threadgroup float4 gk[TB]; + const int h = int(tg.x) / 4, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * 16, w0 = sg * 4; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + device float* sp = state + (i0 * hv + h) * S + j; + const int rs = hv * S; + float4 s0 = float4(sp[0], sp[rs], sp[2 * rs], sp[3 * rs]); + float4 s1 = float4(sp[4 * rs], sp[5 * rs], sp[6 * rs], sp[7 * rs]); + float4 s2 = float4(sp[8 * rs], sp[9 * rs], sp[10 * rs], sp[11 * rs]); + float4 s3 = float4(sp[12 * rs], sp[13 * rs], sp[14 * rs], sp[15 * rs]); + for (int t0 = 0; t0 < T; t0 += TB) { + const int nb = min(TB, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += IB * 32) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; + vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w]; + if (w == 0) { + const float2 g_b = gb[(t0 + tt) * hv + h]; + gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt, slot = t & 1; + const float4 k0 = ks[tt][w0], k1 = ks[tt][w0 + 1], k2 = ks[tt][w0 + 2], k3 = ks[tt][w0 + 3]; + const float4 q0 = qs[tt][w0], q1 = qs[tt][w0 + 1], q2 = qs[tt][w0 + 2], q3 = qs[tt][w0 + 3]; + const float ka = (dot(s0, k0) + dot(s1, k1)) + (dot(s2, k2) + dot(s3, k3)); + const float qa = (dot(s0, q0) + dot(s1, q1)) + (dot(s2, q2) + dot(s3, q3)); + red[slot][sg][lane] = float2(ka, qa); + const float4 g_b = gk[tt]; + const float v = vs[tt][lane]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float2 r = red[slot][0][lane]; + for (int b = 1; b < IB; ++b) r += red[slot][b][lane]; + const float g = g_b.x; + const float delta = (v - g * r.x) * g_b.y; + if (sg == 0) o[(t * hv + h) * S + j] = fma(g, r.y, delta * g_b.z) * scale; + s0 = fma(float4(g), s0, k0 * delta); + s1 = fma(float4(g), s1, k1 * delta); + s2 = fma(float4(g), s2, k2 * delta); + s3 = fma(float4(g), s3, k3 * delta); + } + } + sp[0] = s0.x; sp[rs] = s0.y; sp[2 * rs] = s0.z; sp[3 * rs] = s0.w; + sp[4 * rs] = s1.x; sp[5 * rs] = s1.y; sp[6 * rs] = s1.z; sp[7 * rs] = s1.w; + sp[8 * rs] = s2.x; sp[9 * rs] = s2.y; sp[10 * rs] = s2.z; sp[11 * rs] = s2.w; + sp[12 * rs] = s3.x; sp[13 * rs] = s3.y; sp[14 * rs] = s3.z; sp[15 * rs] = s3.w; +} + +// grid (h_v, T), 128 threads +kernel void gdn_post(device const float* o [[buffer(0)]], device const float* z [[buffer(1)]], + device const float* gamma [[buffer(2)]], device half* yh [[buffer(3)]], device float* y [[buffer(4)]], + constant GdnPArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[4]; + const int h = int(tg.x), t = int(tg.y), hv = int(a.h_v); + const uint i = uint(t * hv + h) * 128 + tid; + const float ov = o[i]; + const float ss = simd_sum(ov * ov); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + const float tot = (part[0] + part[1]) + (part[2] + part[3]); + const float rs = rsqrt(tot / 128.0f + a.eps); + const float zv = gdn_ld_on ? z[uint(t) * uint(GDN_LD) + uint(h) * 128u + tid] : z[i]; + const float yv = ((rs * ov) * gamma[tid]) * (1.0f / (1.0f + exp(-zv))); + yh[i] = half(yv); + if (a.f32) y[i] = yv; +} + +template [[host_name("gdn_scan_1")]] kernel decltype(gdn_scan_t<1, 16>) gdn_scan_t<1, 16>; +template [[host_name("gdn_scan_2")]] kernel decltype(gdn_scan_t<2, 16>) gdn_scan_t<2, 16>; +template [[host_name("gdn_scan_4")]] kernel decltype(gdn_scan_t<4, 16>) gdn_scan_t<4, 16>; +template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 16>) gdn_scan_r<2, 16>; +template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>; +template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 16>) gdn_scan_r<8, 16>; +template [[host_name("gdn_scanr_4_8")]] kernel decltype(gdn_scan_r<4, 8>) gdn_scan_r<4, 8>; +template [[host_name("gdn_scanr_4_4")]] kernel decltype(gdn_scan_r<4, 4>) gdn_scan_r<4, 4>; +template [[host_name("gdn_scanr_4_2")]] kernel decltype(gdn_scan_r<4, 2>) gdn_scan_r<4, 2>; +template [[host_name("gdn_scanr_8_8")]] kernel decltype(gdn_scan_r<8, 8>) gdn_scan_r<8, 8>; +template [[host_name("gdn_scanr_8_4")]] kernel decltype(gdn_scan_r<8, 4>) gdn_scan_r<8, 4>; +template [[host_name("gdn_scanr_2_8")]] kernel decltype(gdn_scan_r<2, 8>) gdn_scan_r<2, 8>; +template [[host_name("gdn_scanr_2_4")]] kernel decltype(gdn_scan_r<2, 4>) gdn_scan_r<2, 4>; + +// ---- gdn_scanf: gdn_scan_r<4, TB> with gdn_post folded in. One threadgroup a value head (grid h_v), 16 simdgroups = +// 4 column blocks x 4 row blocks, 512 threads. The scan is gdn_scan_r's, token for token (the same operations in the +// same order, so o is the same bits); a column block's row-block-0 simdgroup holds o for its 32 columns, writes them and +// their simd_sum of squares to threadgroup memory, and once every TB tokens the whole threadgroup runs gdn_post's +// arithmetic on those tokens (the sum of the four partials in gdn_post's order: the same bits as gdn_post). Outputs: +// y (fp16) and, with `f32`, fp32 y; o itself never reaches device memory. +template +kernel void gdn_scan_f(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]], + device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]], + device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]], + device const float* z [[buffer(6)]], device const float* gamma [[buffer(7)]], + device half* yh [[buffer(8)]], device float* y32 [[buffer(9)]], constant GdnPArgs& a [[buffer(10)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, IB = 4, R = S / IB, R4 = R / 4; + threadgroup float4 qs[TB][32], ks[TB][32]; + threadgroup float2 red[2][16][32]; + threadgroup float vs[TB][128]; + threadgroup float4 gk[TB]; + threadgroup float ot[TB][128]; + threadgroup float osq[TB][4]; + const int h = int(tg.x), hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k); + const int cb = sg >> 2, ib = sg & 3; + const int j = cb * 32 + lane, i0 = ib * R; + const int T = int(a.n_tok); + const float scale = 1.0f / sqrt(float(S)); + float st[R]; +#pragma clang loop unroll(full) + for (int i = 0; i < R; ++i) st[i] = state[((i0 + i) * hv + h) * S + j]; + for (int t0 = 0; t0 < T; t0 += TB) { + const int nb = min(TB, T - t0); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < nb * 32; idx += 512) { + const int tt = idx >> 5, w = idx & 31; + qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w]; + ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w]; + } + for (int idx = tid; idx < nb * 128; idx += 512) { + const int tt = idx >> 7, w = idx & 127; + vs[tt][w] = vc[((t0 + tt) * hv + h) * S + w]; + } + for (int tt = tid; tt < nb; tt += 512) { + const float2 g_b = gb[(t0 + tt) * hv + h]; + gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int tt = 0; tt < nb; ++tt) { + const int t = t0 + tt, slot = t & 1; + float ka = 0.0f, qa = 0.0f; +#pragma clang loop unroll(full) + for (int w = 0; w < R4; ++w) { + const float4 k4 = ks[tt][i0 / 4 + w], q4 = qs[tt][i0 / 4 + w]; + ka = fma(st[4 * w], k4.x, ka); ka = fma(st[4 * w + 1], k4.y, ka); + ka = fma(st[4 * w + 2], k4.z, ka); ka = fma(st[4 * w + 3], k4.w, ka); + qa = fma(st[4 * w], q4.x, qa); qa = fma(st[4 * w + 1], q4.y, qa); + qa = fma(st[4 * w + 2], q4.z, qa); qa = fma(st[4 * w + 3], q4.w, qa); + } + const float4 g_b = gk[tt]; + const float v = vs[tt][j]; + red[slot][sg][lane] = float2(ka, qa); + threadgroup_barrier(mem_flags::mem_threadgroup); + float2 s = red[slot][cb * 4][lane]; +#pragma clang loop unroll(full) + for (int b = 1; b < IB; ++b) s += red[slot][cb * 4 + b][lane]; + const float g = g_b.x; + const float delta = (v - g * s.x) * g_b.y; + if (ib == 0) { + const float ov = fma(g, s.y, delta * g_b.z) * scale; + ot[tt][j] = ov; + const float ss = simd_sum(ov * ov); + if (lane == 0) osq[tt][cb] = ss; + } +#pragma clang loop unroll(full) + for (int w = 0; w < R4; ++w) { + const float4 k4 = ks[tt][i0 / 4 + w]; + st[4 * w] = fma(g, st[4 * w], k4.x * delta); + st[4 * w + 1] = fma(g, st[4 * w + 1], k4.y * delta); + st[4 * w + 2] = fma(g, st[4 * w + 2], k4.z * delta); + st[4 * w + 3] = fma(g, st[4 * w + 3], k4.w * delta); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // gdn_post for these tokens + for (int idx = tid; idx < nb * 128; idx += 512) { + const int tt = idx >> 7, c = idx & 127; + const float tot = (osq[tt][0] + osq[tt][1]) + (osq[tt][2] + osq[tt][3]); + const float rs = rsqrt(tot / 128.0f + a.eps); + const ulong i = ulong((t0 + tt) * hv + h) * 128 + c; + const float ov = ot[tt][c]; + const float yv = ((rs * ov) * gamma[c]) * (1.0f / (1.0f + exp(-z[i]))); + yh[i] = half(yv); + if (a.f32) y32[i] = yv; + } + } +#pragma clang loop unroll(full) + for (int i = 0; i < R; ++i) state[((i0 + i) * hv + h) * S + j] = st[i]; +} +template [[host_name("gdn_scanf_8")]] kernel decltype(gdn_scan_f<8>) gdn_scan_f<8>; +template [[host_name("gdn_scanf_4")]] kernel decltype(gdn_scan_f<4>) gdn_scan_f<4>; +template [[host_name("gdn_scanf_2")]] kernel decltype(gdn_scan_f<2>) gdn_scan_f<2>; + +// ---- the chunked form (chunks of GC = 32 tokens), the WY / UT form of the gated delta rule. With G_t the product +// of the chunk's gates up to t (log space: lg_t) and S0 the state entering the chunk: +// (I + A) D = B, A_ts = beta_t (k_t . k_s) G_t / G_s (s < t), B = [beta v | beta G k] -> D = [U | W] +// delta = U - W S0, o_t = G_t S0^T q_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s, +// S_L = G_L S0 + sum_s G_L / G_s k_s delta_s^T +// (the sequential recurrence's algebra; the rounding differs). Three kernels: +// gdn_cprep per (head, chunk): lg, A, the forward substitution for U and W (a thread per column of B) +// gdn_cstate per (head, 32 state columns): the chunks in order - delta, the inter-chunk output, the state update +// gdn_cout per (head, chunk): o += the intra-chunk part, times 1/sqrt(128) +constant constexpr int GC = 32; + +// grid (h_v, chunks), 256 threads. U, W: [h][chunk][GC][128]; lgo: [h][chunk][GC] +kernel void gdn_cprep(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]], + device const float* vc [[buffer(2)]], device const float2* gb [[buffer(3)]], + device float* U [[buffer(4)]], device float* W [[buffer(5)]], device float* lgo [[buffer(6)]], + constant GdnPArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int S = 128; + threadgroup float A[GC][GC]; + threadgroup float lg[GC], bt[GC]; + const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0); + if (tid == 0) { + float s = 0.0f; + for (int t = 0; t < GC; ++t) { + if (t < nb) s += max(log(gb[(t0 + t) * hv + h].x), -80.0f); + lg[t] = s; + bt[t] = t < nb ? gb[(t0 + t) * hv + h].y : 0.0f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < GC * GC; idx += 256) { + const int t = idx / GC, s = idx % GC; + float v = 0.0f; + if (s < t && t < nb) { + device const float4* kt = (device const float4*)(kn + ((t0 + t) * hkn + hk) * S); + device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S); + float d = 0.0f; + for (int i = 0; i < S / 4; ++i) d += dot(kt[i], ks[i]); + v = bt[t] * d * exp(lg[t] - lg[s]); + } + A[t][s] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // forward substitution, column j of B a thread (threads 0..127: beta v; 128..255: beta G k) + const int j = tid & 127; + const bool isw = tid >= 128; + float x[GC]; + for (int t = 0; t < GC; ++t) { + float b = 0.0f; + if (t < nb) b = isw ? bt[t] * exp(lg[t]) * kn[((t0 + t) * hkn + hk) * S + j] : bt[t] * vc[((t0 + t) * hv + h) * S + j]; + for (int s = 0; s < t; ++s) b = fma(-A[t][s], x[s], b); + x[t] = b; + } + device float* o = (isw ? W : U) + (ulong(h) * NC + c) * GC * S + j; + for (int t = 0; t < GC; ++t) o[t * S] = x[t]; + if (tid < GC) lgo[(h * NC + c) * GC + tid] = lg[tid]; +} + +// grid (h_v * 4), 256 threads = 8 row blocks of 16 x 32 columns; the state's 128 x 32 block in registers (16 a thread) +kernel void gdn_cstate(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]], + device const float* kn [[buffer(2)]], device const float* U [[buffer(3)]], + device const float* W [[buffer(4)]], device const float* lgo [[buffer(5)]], + device float* delta [[buffer(6)]], device float* o [[buffer(7)]], constant GdnPArgs& a [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int S = 128, RB = 16; + threadgroup float red[8][16][32]; + threadgroup float dl[GC][32]; + const int h = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC; + const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * RB; + float st[RB]; + for (int r = 0; r < RB; ++r) st[r] = state[((i0 + r) * hv + h) * S + j]; + for (int c = 0; c < NC; ++c) { + const int t0 = c * GC, nb = min(GC, T - t0); + device const float* Uc = U + (ulong(h) * NC + c) * GC * S; + device const float* Wc = W + (ulong(h) * NC + c) * GC * S; + device const float* lg = lgo + (h * NC + c) * GC; + // pass 0: delta = U - W S0; pass 1: o_inter = G_t S0^T q_t; 16 tokens at a time + for (int pass = 0; pass < 2; ++pass) { + for (int hf = 0; hf < 2; ++hf) { + float p[16]; + for (int u = 0; u < 16; ++u) { + const int t = hf * 16 + u; + float s = 0.0f; + if (t < nb) { + device const float* v = pass == 0 ? Wc + t * S + i0 : qn + ((t0 + t) * hkn + hk) * S + i0; + for (int r = 0; r < RB; ++r) s = fma(v[r], st[r], s); + } + p[u] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int u = 0; u < 16; ++u) red[sg][u][lane] = p[u]; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int u = sg * 2; u < sg * 2 + 2; ++u) { + const int t = hf * 16 + u; + float s = 0.0f; + for (int b = 0; b < 8; ++b) s += red[b][u][lane]; + if (t < nb) { + if (pass == 0) { + const float d = Uc[t * S + j] - s; + dl[t][lane] = d; + delta[((ulong(h) * NC + c) * GC + t) * S + j] = d; + } else { + o[((t0 + t) * hv + h) * S + j] = exp(lg[t]) * s; + } + } else if (pass == 0) { + dl[t][lane] = 0.0f; + } + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // S = G_L S0 + sum_s G_L / G_s k_s delta_s^T + const float lgl = lg[nb - 1], gl = exp(lgl); + for (int r = 0; r < RB; ++r) st[r] *= gl; + for (int s = 0; s < nb; ++s) { + const float f = exp(lgl - lg[s]) * dl[s][lane]; + device const float* ks = kn + ((t0 + s) * hkn + hk) * S + i0; + for (int r = 0; r < RB; ++r) st[r] = fma(ks[r], f, st[r]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (int r = 0; r < RB; ++r) state[((i0 + r) * hv + h) * S + j] = st[r]; +} + +// grid (h_v, chunks), 128 threads (column j): o_t = (o_inter_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s) / sqrt(128) +kernel void gdn_cout(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]], + device const float* delta [[buffer(2)]], device const float* lgo [[buffer(3)]], + device float* o [[buffer(4)]], constant GdnPArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int S = 128; + threadgroup float M[GC][GC]; + const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0); + device const float* lg = lgo + (h * NC + c) * GC; + for (int idx = tid; idx < GC * GC; idx += 128) { + const int t = idx / GC, s = idx % GC; + float v = 0.0f; + if (s <= t && t < nb) { + device const float4* qt = (device const float4*)(qn + ((t0 + t) * hkn + hk) * S); + device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S); + float d = 0.0f; + for (int i = 0; i < S / 4; ++i) d += dot(qt[i], ks[i]); + v = d * exp(lg[t] - lg[s]); + } + M[t][s] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int j = tid; + device const float* dc = delta + (ulong(h) * NC + c) * GC * S + j; + float dv[GC]; + for (int s = 0; s < GC; ++s) dv[s] = s < nb ? dc[s * S] : 0.0f; + const float scale = 1.0f / sqrt(float(S)); + for (int t = 0; t < nb; ++t) { + float s = o[((t0 + t) * hv + h) * S + j]; + for (int u = 0; u <= t; ++u) s = fma(M[t][u], dv[u], s); + o[((t0 + t) * hv + h) * S + j] = s * scale; + } +} + +// ---- the chunked form v2: the same algebra as gdn_cprep / gdn_cstate / gdn_cout (WY / UT, chunks of GC = 32), organised for the +// GPU. Staging the chunk's k in threadgroup memory removes the redundant device loads of gdn_cprep, the state update, the +// inter-chunk and the intra-chunk products run as 8x8 simdgroup_matrix products (fp32), and the intra-chunk part is folded into +// the state kernel (no gdn_cout, no delta buffer). The chunk's q.k matrix M (with the decay factors) comes from the prep. +// gdn_cprep2 per (head, chunk): lg, A, the forward substitution for U and W, and M_ts = (q_t . k_s) G_t / G_s (s <= t) +// gdn_cstate2 per (head, 32 state columns): the chunks in order - delta = U - W S, o = (G q S + M delta) / sqrt(128), +// S = G_L S + sum_s (G_L / G_s) k_s delta_s^T. The q and k rows of the last chunk's 32 tokens are read in +// full (rows past n_tok are read and multiplied by zero): the caller's buffers hold a multiple of 32 rows. +kernel void gdn_cprep2(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]], + device const float* vc [[buffer(2)]], device const float2* gb [[buffer(3)]], + device float* U [[buffer(4)]], device float* W [[buffer(5)]], device float* lgo [[buffer(6)]], + device float* Mo [[buffer(7)]], constant GdnPArgs& a [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int S = 128; + threadgroup float4 Ks[GC * 32]; + threadgroup float4 A4[GC * 8]; + threadgroup float* A = (threadgroup float*)A4; + threadgroup float lg[GC], bt[GC]; + const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0); + if (tid == 0) { + float s = 0.0f; + for (int t = 0; t < GC; ++t) { + if (t < nb) s += max(log(gb[(t0 + t) * hv + h].x), -80.0f); + lg[t] = s; + bt[t] = t < nb ? gb[(t0 + t) * hv + h].y : 0.0f; + } + } + device const float4* kn4 = (device const float4*)kn; + for (int idx = tid; idx < GC * 32; idx += 256) { + const int t = idx >> 5, w = idx & 31; + Ks[idx] = t < nb ? kn4[((t0 + t) * hkn + hk) * 32 + w] : float4(0.0f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < GC * GC; idx += 256) { + const int t = idx / GC, s = idx % GC; + float av = 0.0f, mv = 0.0f; + if (t < nb && s <= t && gdn_probe != 1) { + device const float4* qt = (device const float4*)(qn + ((t0 + t) * hkn + hk) * S); + float dm = 0.0f, da = 0.0f; + for (int i = 0; i < S / 4; ++i) { + const float4 ks = Ks[s * 32 + i]; + dm += dot(qt[i], ks); + da += dot(Ks[t * 32 + i], ks); + } + const float f = exp(lg[t] - lg[s]); + mv = dm * f; + if (s < t) av = bt[t] * da * f; + } + A[t * GC + s] = av; + Mo[((ulong(h) * NC + c) * GC + t) * GC + s] = mv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int j = tid & 127; + const bool isw = tid >= 128; + threadgroup const float* Kf = (threadgroup const float*)Ks; + float x[GC]; +#pragma clang loop unroll(full) + for (int t = 0; t < GC; ++t) x[t] = 0.0f; +#pragma clang loop unroll(full) + for (int t = 0; t < GC; ++t) { + float b = 0.0f; + if (t < nb) b = isw ? bt[t] * exp(lg[t]) * Kf[t * S + j] : bt[t] * vc[((t0 + t) * hv + h) * S + j]; +#pragma clang loop unroll(full) + for (int s4 = 0; s4 < (t + 3) / 4; ++s4) { + const float4 ar = A4[t * 8 + s4]; + b = fma(-ar.x, x[4 * s4], b); + b = fma(-ar.y, x[4 * s4 + 1], b); + b = fma(-ar.z, x[4 * s4 + 2], b); + b = fma(-ar.w, x[4 * s4 + 3], b); + } + x[t] = b; + } + device float* o = (isw ? W : U) + (ulong(h) * NC + c) * GC * S + j; +#pragma clang loop unroll(full) + for (int t = 0; t < GC; ++t) o[t * S] = x[t]; + if (tid < GC) lgo[(h * NC + c) * GC + tid] = lg[tid]; +} + + +// grid (h_v * 4), 128 threads = 4 simdgroups. A simdgroup owns rows 32 sg .. 32 sg + 31 of the 128 x 32 state block as 4 x 4 +// simdgroup matrices that stay in registers for the whole call; each chunk the block is stored to threadgroup memory (Sm) once, +// where the two products that need all 128 rows (W S and Q S, sharing the S tiles) read it. +kernel void gdn_cstate2(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]], + device const float* kn [[buffer(2)]], device const float* U [[buffer(3)]], + device const float* W [[buffer(4)]], device const float* lgo [[buffer(5)]], + device const float* Mo [[buffer(6)]], device float* o [[buffer(7)]], constant GdnPArgs& a [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr int S = 128; + threadgroup float Sm[S][32]; + threadgroup float Dl[GC][32]; + threadgroup float X[GC][32]; + threadgroup float Y[GC][32]; + const int h = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok); + const int NC = (T + GC - 1) / GC, j0 = (int(tg.x) % 4) * 32; + const float scale = 1.0f / sqrt(float(S)); + simdgroup_float8x8 st[4][4]; + for (int rt = 0; rt < 4; ++rt) + for (int jt = 0; jt < 4; ++jt) + simdgroup_load(st[rt][jt], state + ((32 * sg + 8 * rt) * hv + h) * S + j0 + 8 * jt, hv * S); + for (int c = 0; c < NC; ++c) { + const int t0 = c * GC, nb = min(GC, T - t0); + device const float* Uc = U + (ulong(h) * NC + c) * GC * S; + device const float* Wc = W + (ulong(h) * NC + c) * GC * S; + device const float* Mc = Mo + (ulong(h) * NC + c) * GC * GC; + device const float* lg = lgo + (h * NC + c) * GC; + for (int rt = 0; rt < 4; ++rt) + for (int jt = 0; jt < 4; ++jt) simdgroup_store(st[rt][jt], &Sm[32 * sg + 8 * rt][8 * jt], 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + // P1 = W S0 (into X) and P2 = Q S0 (into Y), the S tiles read once for both + simdgroup_float8x8 acc1[4], acc2[4]; + for (int jt = 0; jt < 4; ++jt) { + acc1[jt] = make_filled_simdgroup_matrix(0.0f); + acc2[jt] = make_filled_simdgroup_matrix(0.0f); + } + device const float* qrow = qn + ((t0 + 8 * sg) * hkn + hk) * S; + if (gdn_probe != 1) +#pragma clang loop unroll(full) + for (int kb = 0; kb < 16; ++kb) { + simdgroup_float8x8 aw, aq, bs; + simdgroup_load(aw, Wc + (8 * sg) * S + 8 * kb, S); + simdgroup_load(aq, qrow + 8 * kb, hkn * S); + for (int jt = 0; jt < 4; ++jt) { + simdgroup_load(bs, &Sm[8 * kb][8 * jt], 32); + simdgroup_multiply_accumulate(acc1[jt], aw, bs, acc1[jt]); + simdgroup_multiply_accumulate(acc2[jt], aq, bs, acc2[jt]); + } + } + for (int jt = 0; jt < 4; ++jt) { + simdgroup_store(acc1[jt], &X[8 * sg][8 * jt], 32); + simdgroup_store(acc2[jt], &Y[8 * sg][8 * jt], 32); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < GC * 32; idx += 128) { + const int t = idx >> 5, j = idx & 31; + Dl[t][j] = t < nb ? Uc[t * S + j0 + j] - X[t][j] : 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // intra = M delta (into X) + for (int jt = 0; jt < 4; ++jt) acc1[jt] = make_filled_simdgroup_matrix(0.0f); + if (gdn_probe != 2) +#pragma clang loop unroll(full) + for (int kb = 0; kb < 4; ++kb) { + simdgroup_float8x8 am, bd; + simdgroup_load(am, Mc + (8 * sg) * GC + 8 * kb, GC); + for (int jt = 0; jt < 4; ++jt) { + simdgroup_load(bd, &Dl[8 * kb][8 * jt], 32); + simdgroup_multiply_accumulate(acc1[jt], am, bd, acc1[jt]); + } + } + for (int jt = 0; jt < 4; ++jt) simdgroup_store(acc1[jt], &X[8 * sg][8 * jt], 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const float lgl = lg[nb - 1], gl = exp(lgl); + if (gdn_probe != 4) + for (int idx = tid; idx < GC * 32; idx += 128) { + const int t = idx >> 5, j = idx & 31; + if (t < nb) { + o[((t0 + t) * hv + h) * S + j0 + j] = fma(exp(lg[t]), Y[t][j], X[t][j]) * scale; + Dl[t][j] *= exp(lgl - lg[t]); + } else { + Dl[t][j] = 0.0f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // S = G_L S + K^T delta' + for (int rt = 0; rt < 4; ++rt) + for (int jt = 0; jt < 4; ++jt) { + st[rt][jt].thread_elements()[0] *= gl; + st[rt][jt].thread_elements()[1] *= gl; + } + if (gdn_probe != 3) +#pragma clang loop unroll(full) + for (int kb = 0; kb < 4; ++kb) { + simdgroup_float8x8 bd[4]; + for (int jt = 0; jt < 4; ++jt) simdgroup_load(bd[jt], &Dl[8 * kb][8 * jt], 32); + for (int rt = 0; rt < 4; ++rt) { + simdgroup_float8x8 ak; + simdgroup_load(ak, kn + ((t0 + 8 * kb) * hkn + hk) * S + 32 * sg + 8 * rt, hkn * S, ulong2(0, 0), true); + for (int jt = 0; jt < 4; ++jt) simdgroup_multiply_accumulate(st[rt][jt], ak, bd[jt], st[rt][jt]); + } + } + } + for (int rt = 0; rt < 4; ++rt) + for (int jt = 0; jt < 4; ++jt) + simdgroup_store(st[rt][jt], state + ((32 * sg + 8 * rt) * hv + h) * S + j0 + 8 * jt, hv * S); +} diff --git a/edge0-strata/metal/kernels/gdn_tprobe.metal b/edge0-strata/metal/kernels/gdn_tprobe.metal new file mode 100644 index 00000000..841609c4 --- /dev/null +++ b/edge0-strata/metal/kernels/gdn_tprobe.metal @@ -0,0 +1,75 @@ +// metal/kernels/gdn_tprobe.metal - go / no-go probe for a tensor-op (MPP matmul2d) chunk-state kernel: the five products of one GDN +// chunk step (W S, Q S, M delta, K^T delta' accumulated into S) on fp16 operands, S kept in fp32 threadgroup memory, 64 chunk steps +// a threadgroup, grid h_v * 4 (the shape of gdn_cstate2). No numerics: timing only. +#include +using namespace metal; +using namespace mpp::tensor_ops; + +constant int PROBE [[function_constant(0)]]; +constant int probe = is_function_constant_defined(PROBE) ? PROBE : 0; // 1 = skip the products (SIMT parts only) + +kernel void gdn_tprobe(device float* out [[buffer(0)]], device half* Wh [[buffer(1)]], device half* Qh [[buffer(2)]], + device half* Kh [[buffer(3)]], device half* Mh [[buffer(4)]], constant int& NC [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup half Sh[128 * 32]; + threadgroup half Dh[32 * 32]; + threadgroup float X[32 * 32], Y[32 * 32], Z[32 * 32]; + tensor, tensor_inline> tSh(Sh, dextents(32, 128)); + tensor, tensor_inline> tDh(Dh, dextents(32, 32)); + tensor, tensor_inline> tX(X, dextents(32, 32)); + tensor, tensor_inline> tY(Y, dextents(32, 32)); + tensor, tensor_inline> tZ(Z, dextents(32, 32)); + constexpr auto d_ws = matmul2d_descriptor(32, 32, 128, false, false, false); // [32 x 128] . [128 x 32] + constexpr auto d_md = matmul2d_descriptor(32, 32, 32, false, false, false); // [32 x 32] . [32 x 32] + constexpr auto d_up = matmul2d_descriptor(128, 32, 32, true, false, false, matmul2d_descriptor::mode::multiply_accumulate); // K^T delta' + matmul2d> op_ws; + matmul2d> op_md; + matmul2d> op_up; + const ulong hb = ulong(tg.x) / 4; + using KT = tensor, tensor_inline>; + using DT = tensor, tensor_inline>; + auto cS = op_up.get_destination_cooperative_tensor(); +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < cS.get_capacity(); ++i) + if (cS.is_valid_element(i)) cS[i] = 0.001f * float(i & 15); + for (int c = 0; c < NC; ++c) { + const ulong cb = (hb * NC + c); +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < cS.get_capacity(); ++i) + if (cS.is_valid_element(i)) { + auto ids = cS.get_multidimensional_index(i); + Sh[ids[1] * 32 + ids[0]] = half(cS[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (probe != 1) { + tensor, tensor_inline> tW(Wh + cb * 32 * 128, dextents(128, 32)); + tensor, tensor_inline> tQ(Qh + cb * 32 * 128, dextents(128, 32)); + op_ws.run(tW, tSh, tX); + op_ws.run(tQ, tSh, tY); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int i = tid; i < 32 * 32; i += 128) Dh[i] = half(0.01f * X[i]); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (probe != 1) { + tensor, tensor_inline> tM(Mh + cb * 32 * 32, dextents(32, 32)); + op_md.run(tM, tDh, tZ); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int i = tid; i < 32 * 32; i += 128) Dh[i] = half(Z[i] + 0.01f * Y[i]); +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < cS.get_capacity(); ++i) + if (cS.is_valid_element(i)) cS[i] *= 0.99f; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (probe != 1) { + tensor, tensor_inline> tK(Kh + cb * 32 * 128, dextents(128, 32)); + op_up.run(tK, tDh, cS); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < cS.get_capacity(); ++i) + if (cS.is_valid_element(i)) { + auto ids = cS.get_multidimensional_index(i); + out[ulong(tg.x) * 4096 + ids[1] * 32 + ids[0]] = cS[i]; + } +} diff --git a/edge0-strata/metal/kernels/gemv.metal b/edge0-strata/metal/kernels/gemv.metal new file mode 100644 index 00000000..6d4ed3a5 --- /dev/null +++ b/edge0-strata/metal/kernels/gemv.metal @@ -0,0 +1,264 @@ +// metal/kernels/gemv.metal - y[t][r] = W[r] . x[t] for a window of 1..8 tokens (dense projections, the shared +// expert, the LM head), any format of quant.metal/units.metal. +// +// Shape of a threadgroup: SG simdgroups; KS of them split the row's units (KS = 1: no split), the other SG/KS cover +// different rows. In a simdgroup, LPR lanes share a row and 32/LPR rows run side by side; each lane holds NR rows. +// The weights of a unit are decoded once and applied to every token. Partial sums of a K split are added in a +// fixed order (deterministic). + +struct GemvArgs { + uint n_in; + uint n_out; + uint n_tok; + uint row_bytes; + uint x_stride; // floats between tokens' x rows + uint y_stride; // floats between tokens' y rows +}; + +constant int NTOK [[function_constant(0)]]; // tokens of the window (1..8): the token loops are specialised +// a bench probe (function constant 9): 1 = tokens past the first reuse token 0's x (no load; wrong values), +// 2 = tokens past the first skip the combine (wrong values) +constant int GEMV_PROBE [[function_constant(9)]]; +constant int gemv_probe = is_function_constant_defined(GEMV_PROBE) ? GEMV_PROBE : 0; + +template +kernel void gemv(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RS = SG / KS, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + threadgroup float red[KS > 1 ? KS * RS * RG * NR * 8 : 1]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int ks = sg / RS, rs = sg % RS; + const int li = lane % LPR; + const int lrow0 = (rs * RG + lane / LPR) * NR; + const int row0 = int(tg.x) * RS * RG * NR + lrow0; + const int U = int(a.n_in) / UNIT; + const int u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + float acc[NR][8]; + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; + for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; + const int s16n = int(a.n_in) / 16; + if (gemv_probe == 3) { + // rows outer, tokens inner: one row's decoded unit in registers at a time (x re-read per row from L1) + for (int u = u0 + li; u < u1; u += LPR) + for (int r = 0; r < NR; ++r) { + float q[UNIT]; + float4 sc; + Dec::run(rp[r], u, q, sc, tab); + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + acc[r][t] += unit_combine(q, sc, xs, xsum); + } + } + } else + for (int u = u0 + li; u < u1; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); + float xs[UNIT], xsum[2]; + for (int t = 0; t < NTOK; ++t) { + if (gemv_probe != 1 || t == 0) X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + if (gemv_probe != 2 || t == 0) + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + constexpr int RPT = RS * RG * NR; + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (KS == 1) { + if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; + } else if (li == 0) { + red[(ks * RPT + lrow0 + r) * 8 + t] = v; + } + } + if (KS > 1) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && li == 0) + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NTOK; ++t) { + float v = 0.0f; + for (int k = 0; k < KS; ++k) v += red[(k * RPT + lrow0 + r) * 8 + t]; + if (row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; + } + } +} + +#define GV(F, L, R, S, K) template [[host_name("gemv_" #F "_" #L "_" #R "_" #S "_" #K)]] kernel decltype(gemv) gemv; +#define GV_SET(F) GV(F, 32, 1, 4, 1) GV(F, 32, 2, 4, 1) GV(F, 16, 1, 4, 1) GV(F, 16, 2, 4, 1) GV(F, 8, 1, 4, 1) \ + GV(F, 8, 2, 4, 1) GV(F, 4, 2, 4, 1) GV(F, 4, 1, 4, 1) GV(F, 32, 1, 4, 4) GV(F, 32, 2, 4, 2) \ + GV(F, 32, 1, 8, 8) GV(F, 16, 1, 8, 4) GV(F, 32, 1, 8, 2) GV(F, 8, 1, 8, 1) GV(F, 1, 1, 4, 1) GV(F, 1, 2, 4, 1) GV(F, 2, 1, 4, 1) GV(F, 2, 2, 4, 1) GV(F, 1, 1, 2, 1) GV(F, 4, 4, 4, 1) +#ifndef GEMV_TYPES +#define GEMV_TYPES GV_SET(T_F32) GV_SET(T_BF16) GV_SET(T_Q8_0) GV_SET(T_Q4_0) GV_SET(T_Q5_0) GV_SET(T_IQ4_NL) GV_SET(T_IQ4_XS) \ + GV_SET(T_Q2_0) GV_SET(T_IQ3_S) GV_SET(T_Q3_K) GV_SET(T_Q4_K) GV_SET(T_Q5_K) GV_SET(T_Q6_K) +#endif +GEMV_TYPES + +// The shared expert's gate and up in one dispatch: h[t][r] = silu(Wg[r] . x[t]) * (Wu[r] . x[t]), each with its own +// format (and so its own prepared activation kind); h is written raw (its consumer reads raw activations). +// KS = 1; tiles as gemv. Buffers: Wg, xg (Wg's kind), Wu, xu (Wu's kind), s16, h, args (row_bytes = Wg's), rb_up. +template +kernel void gemv_pair(device const uchar* Wg [[buffer(0)]], device const float* xg [[buffer(1)]], + device const uchar* Wu [[buffer(2)]], device const float* xu [[buffer(3)]], + device const float* s16 [[buffer(4)]], device float* h [[buffer(5)]], + constant GemvArgs& a [[buffer(6)]], constant uint& rb_up [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int U1 = fmt::UNIT, U2 = fmt::UNIT, RG = 32 / LPR; + constexpr int TB = fmt::TAB > fmt::TAB ? fmt::TAB : fmt::TAB; + threadgroup uint tab1[fmt::TAB]; + threadgroup uint tab2[fmt::TAB]; + (void)TB; + load_tab(tab1, tid, SG * 32); + load_tab(tab2, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int li = lane % LPR; + const int row = (int(tg.x) * SG + sg) * RG + lane / LPR; + const int rr = min(row, int(a.n_out) - 1); + device const uchar* rg = Wg + ulong(rr) * a.row_bytes; + device const uchar* ru = Wu + ulong(rr) * rb_up; + const int s16n = int(a.n_in) / 16; + float ag[8], au[8]; + for (int t = 0; t < 8; ++t) ag[t] = au[t] = 0.0f; + for (int u = li; u < int(a.n_in) / U1; u += LPR) { + float q[U1]; + float4 sc; + Dec::run(rg, u, q, sc, tab1); + for (int t = 0; t < NTOK; ++t) { + float xs[U1], xsum[2]; + X::load(xg + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + ag[t] += unit_combine(q, sc, xs, xsum); + } + } + for (int u = li; u < int(a.n_in) / U2; u += LPR) { + float q[U2]; + float4 sc; + Dec::run(ru, u, q, sc, tab2); + for (int t = 0; t < NTOK; ++t) { + float xs[U2], xsum[2]; + X::load(xu + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + au[t] += unit_combine(q, sc, xs, xsum); + } + } + for (int t = 0; t < NTOK; ++t) { + const float gv = lane_sum(ag[t]), uv = lane_sum(au[t]); + if (li == 0 && row < int(a.n_out)) h[t * int(a.y_stride) + row] = gv / (1.0f + exp(-gv)) * uv; + } +} + +#define GP_(F1, F2) template [[host_name("gemv_pair_" #F1 "_" #F2)]] kernel decltype(gemv_pair) gemv_pair; +#define GP(F1, F2) GP_(F1, F2) +#define GP_ROW(F1) GP(F1, T_IQ3_S) GP(F1, T_IQ4_XS) GP(F1, T_Q2_0) GP(F1, T_Q3_K) GP(F1, T_Q4_K) GP(F1, T_Q5_K) GP(F1, T_Q6_K) +GP_ROW(T_IQ3_S) GP_ROW(T_IQ4_XS) GP_ROW(T_Q2_0) GP_ROW(T_Q3_K) GP_ROW(T_Q4_K) GP_ROW(T_Q5_K) GP_ROW(T_Q6_K) +// Q2_0 with raw activations (the shared expert's down reading the pair's raw h) +GV_SET(T_Q2_0R) + +// gemv_cat: two matrices of the same format and n_in over the same x in one dispatch - rows [0, n_split) are W's +// (into y), rows [n_split, n_out) are W2's (into y2, from its row 0). Each row's arithmetic is gemv's (KS = 1), so the +// outputs are bitwise the two separate GEMVs'. (GDN's ssm_alpha + ssm_beta: 48 + 48 rows.) +// Buffers: W, xp, s16, y, args (n_out = both), W2, y2, n_split. +template +kernel void gemv_cat(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + device const uchar* W2 [[buffer(5)]], device float* y2 [[buffer(6)]], constant uint& n_split [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int li = lane % LPR; + const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; + const int U = int(a.n_in) / UNIT, sp = int(n_split); + float acc[NR][8]; + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; + for (int r = 0; r < NR; ++r) { + const int row = min(row0 + r, int(a.n_out) - 1); + rp[r] = row < sp ? W + (ulong)row * a.row_bytes : W2 + (ulong)(row - sp) * a.row_bytes; + } + const int s16n = int(a.n_in) / 16; + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(a.x_stride), s16 + t * s16n, u, xs, xsum); + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + const int row = row0 + r; + if (li == 0 && row < int(a.n_out)) { + if (row < sp) y[t * int(a.y_stride) + row] = v; + else y2[t * int(a.y_stride) + row - sp] = v; + } + } +} +#define GVC_(F, L, R, S) template [[host_name("gemvc_" #F "_" #L "_" #R "_" #S)]] kernel decltype(gemv_cat) gemv_cat; +#define GVC(F, L, R, S) GVC_(F, L, R, S) +GVC(T_BF16, 32, 1, 4) GVC(T_BF16, 32, 1, 8) GVC(T_BF16, 16, 1, 4) GVC(T_BF16, 32, 2, 4) +GVC(T_F32, 32, 1, 4) GVC(T_F32, 32, 1, 8) GVC(T_F32, 16, 1, 4) GVC(T_F32, 32, 2, 4) + +// gemv_x: gemv (KS = 1) with the window's prepared x (the format's plane, NTOK rows, and its per-16 sums) staged in +// threadgroup memory once per threadgroup, so the per-unit x reads of every token come from threadgroup memory +// instead of L1 (on the M4 the x traffic through L1 - 4 B a weight a token, shared by the rows in a simdgroup - is +// what a window token costs). Same tile meaning and the same arithmetic as gemv_____1: bitwise. +// Dynamic threadgroup memory 0: 4 * NTOK * (n_in + n_in / 16) bytes. +template +kernel void gemv_x(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* y [[buffer(3)]], constant GemvArgs& a [[buffer(4)]], + threadgroup float* xt [[threadgroup(0)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + const int n = int(a.n_in), s16n = n / 16; + threadgroup float* st = xt + NTOK * n; + for (int t = 0; t < NTOK; ++t) { + device const float4* src = (device const float4*)(xp + t * int(a.x_stride)); + threadgroup float4* dst = (threadgroup float4*)(xt + t * n); + for (int i = tid; i < n / 4; i += SG * 32) dst[i] = src[i]; + for (int i = tid; i < s16n; i += SG * 32) st[t * s16n + i] = s16[t * s16n + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int li = lane % LPR; + const int row0 = int(tg.x) * SG * RG * NR + (sg * RG + lane / LPR) * NR; + const int U = n / UNIT; + float acc[NR][8]; + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NTOK; ++t) acc[r][t] = 0.0f; + device const uchar* rp[NR]; + for (int r = 0; r < NR; ++r) rp[r] = W + (ulong)min(row0 + r, int(a.n_out) - 1) * a.row_bytes; + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; + for (int r = 0; r < NR; ++r) Dec::run(rp[r], u, q[r], sc[r], tab); + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load((threadgroup const float*)(xt + t * n), (threadgroup const float*)(st + t * s16n), u, xs, xsum); + for (int r = 0; r < NR; ++r) acc[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum(acc[r][t]); + if (li == 0 && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = v; + } +} +#define GVX(F, L, R, S) template [[host_name("gemvx_" #F "_" #L "_" #R "_" #S)]] kernel decltype(gemv_x) gemv_x; +#define GVX_SET(F) GVX(F, 32, 1, 4) GVX(F, 32, 2, 4) GVX(F, 16, 1, 4) GVX(F, 16, 2, 4) GVX(F, 8, 1, 4) GVX(F, 8, 2, 4) \ + GVX(F, 4, 2, 4) GVX(F, 16, 2, 8) GVX(F, 8, 2, 8) GVX(F, 16, 4, 4) GVX(F, 32, 2, 8) GVX(F, 2, 2, 4) +GVX_SET(T_BF16) GVX_SET(T_Q8_0) GVX_SET(T_Q4_0) GVX_SET(T_Q5_0) GVX_SET(T_IQ4_NL) GVX_SET(T_IQ4_XS) GVX_SET(T_Q2_0) +GVX_SET(T_IQ3_S) GVX_SET(T_Q3_K) GVX_SET(T_Q4_K) GVX_SET(T_Q5_K) GVX_SET(T_Q6_K) diff --git a/edge0-strata/metal/kernels/glue.metal b/edge0-strata/metal/kernels/glue.metal new file mode 100644 index 00000000..b7bc8bcc --- /dev/null +++ b/edge0-strata/metal/kernels/glue.metal @@ -0,0 +1,63 @@ +// metal/kernels/glue.metal - small elementwise steps between the big kernels (after quant/units/moe.metal). +// +// swiglu_prep h = silu(g) * u for the shared expert, written as prepared activations (xprep layout) +// moe_combine_sh bo[t] = sum_k ye[entry(t, k)] + sigmoid(sg[t]) * ysh[t] (k in order) + +struct GlueArgs { + uint n; // vector width + uint n_tok; + uint n_embd; + uint topk; +}; + +// grid (n / 16, n_tok), 16 threads: one group of 16 values +kernel void swiglu_prep(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], + device float* xp [[buffer(2)]], constant GlueArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort j [[thread_index_in_threadgroup]]) { + threadgroup float v[16]; + const uint t = tg.y, d = tg.x * 16 + j, plane = a.n_tok * a.n; + const float gv = g[t * a.n + d]; + const float h = gv / (1.0f + exp(-gv)) * u[t * a.n + d]; + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; + const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + device float* raw = xp + t * a.n + d; + raw[0] = h; + raw[plane] = h * q2s[j & 7]; + raw[2 * plane] = h * p4s[j & 3]; + raw[3 * plane] = (j & 1) ? h * 0.00390625f : h; + v[j] = h; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (j == 0) { + float s0 = 0.0f, s1 = 0.0f; + for (int k = 0; k < 16; k += 2) { s0 += v[k]; s1 += v[k + 1]; } + xp[4 * plane + t * (a.n / 16) + tg.x] = s0 + s1; + } +} + +kernel void moe_combine_sh(device const float* ye [[buffer(0)]], device const int* plan [[buffer(1)]], + device const float* ysh [[buffer(2)]], device const float* sg [[buffer(3)]], + device float* bo [[buffer(4)]], constant GlueArgs& a [[buffer(5)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; + float s = 0.0f; + for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * a.topk + j] * int(a.n_embd) + k]; + const float gs = 1.0f / (1.0f + exp(-sg[t])); + bo[gid] = fma(gs, ysh[gid], s); +} + +// slc_touch: reads n16 16-byte words (a weight range the next kernels will need) so they sit in the SLC; the xor is +// written only if it equals a value it never takes (keeps the loads alive). grid (ceil(n16 / 1024)), 256 threads, +// 4 words a thread. +kernel void slc_touch(device const uint4* src [[buffer(0)]], device uint* sink [[buffer(1)]], constant uint& n16 [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + uint s = 0; + for (uint k = 0; k < 4; ++k) { + const uint j = (gid >> 8) * 1024 + k * 256 + (gid & 255); + if (j < n16) { + const uint4 v = src[j]; + s ^= v.x ^ v.y ^ v.z ^ v.w; + } + } + if (s == 0x9e3779b9u) sink[0] = s; +} diff --git a/edge0-strata/metal/kernels/gr.metal b/edge0-strata/metal/kernels/gr.metal new file mode 100644 index 00000000..c1960fd5 --- /dev/null +++ b/edge0-strata/metal/kernels/gr.metal @@ -0,0 +1,179 @@ +// metal/kernels/gr.metal - the hyper-connection (gated residual) read for a window of tokens, with the previous +// sublayer's write folded in. The math is fused_gr.hpp's (native contract: FP32 activations, BF16 weights): +// +// gr_norm R'[c] = R[c] + bo * 2 sigmoid(inj_prev[c] / hc) (when apply; written back in place) +// xn[c] = R'[c] * w_norm[c] * rsqrt(mean(R'[c]^2) + eps) +// gr_down lo[k] = silu((w_down[k] . xn) / hc) k < hc_lr; inject[c] = w_inject[c] . xn +// gr_up mixed[d] = mean_c xn[c][d] * sigmoid(w_up[c n + d] . lo), written as prepared activations (xprep's +// planes and per-16 sums) for the GEMVs that read it +// +// Shapes: n_embd 2560, hc 4, hc_lr 320. Per token t: R[t] (hc * n), bo[t] (n), inj_prev[t] (hc), xn[t] (hc * n), +// lo[t] (hc_lr), inject[t] (hc), XP planes of n_tok rows. + +constant int GR_NTOK [[function_constant(1)]]; + +struct GrArgs { + uint n_embd; + uint hc; + uint hc_lr; + uint n_tok; + uint apply; + float eps; +}; + +// one threadgroup per (stream c, token t), 256 threads +kernel void gr_norm(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + float ss = 0.0f; + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + v = fma(bo[t * n + d], g, v); + r[d] = v; + } + ss = fma(v, v, ss); + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + for (uint d = tid; d < n; d += 256) o[d] = r[d] * w[d] * rs; +} + +// rows: 0 .. hc_lr - 1 = w_down, hc_lr .. hc_lr + hc - 1 = w_inject (both bf16 rows of hc * n_embd). +// SG simdgroups split the row's units KS ways; LPR = 32 (a simdgroup per row); RS = SG / KS rows per threadgroup. +// one row's K-split partial over xn (units of format F) +template +static inline void grd_acc(device const uchar* w, int nin, int ks, int KS, ushort lane, device const float* xn, + thread float* acc, threadgroup const uint* notab) { + constexpr int UNIT = fmt::UNIT; + const int U = nin / UNIT, u0 = ks * U / KS, u1 = (ks + 1) * U / KS; + for (int u = u0 + lane; u < u1; u += 32) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, notab); + for (int t = 0; t < GR_NTOK; ++t) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + for (int k = 0; k < UNIT; k += 16) load16(xn + t * nin + UNIT * u + k, xs + k); + acc[t] += unit_combine(q, sc, xs, xsum); + } + } +} + +// FI: the inject rows' format (the Unsloth files: Q8_0 down, F32 inject); FI = F is the original kernel +template +kernel void gr_down_f(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* xn [[buffer(2)]], device float* lo [[buffer(3)]], + device float* inject [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + constexpr int RS = SG / KS; + threadgroup float red[KS * RS * 8]; + threadgroup uint notab[1]; + const int ks = sg / RS, rs = sg % RS; + const int row = int(tg.x) * RS + rs; + const int nrows = int(a.hc_lr + a.hc); + const int rr = min(row, nrows - 1); + const int nin = int(a.hc * a.n_embd); + const uint rb = uint(nin / fmt::QK * fmt::BB), rbi = uint(nin / fmt::QK * fmt::BB); + float acc[8]; + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + if (rr < int(a.hc_lr)) grd_acc(w_down + ulong(rr) * rb, nin, ks, KS, lane, xn, acc, notab); + else grd_acc(w_inject + ulong(rr - int(a.hc_lr)) * rbi, nin, ks, KS, lane, xn, acc, notab); + for (int t = 0; t < GR_NTOK; ++t) { + const float v = simd_sum(acc[t]); + if (lane == 0) red[(ks * RS + rs) * 8 + t] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0 && lane < GR_NTOK && row < nrows) { + const int t = lane; + float v = 0.0f; + for (int k = 0; k < KS; ++k) v += red[(k * RS + rs) * 8 + t]; + if (row < int(a.hc_lr)) { + const float z = v / float(a.hc); + lo[t * int(a.hc_lr) + row] = z / (1.0f + exp(-z)); + } else { + inject[t * int(a.hc) + row - int(a.hc_lr)] = v; + } + } +} + +// A threadgroup owns 16 consecutive d for all hc = 4 streams: 8 simdgroups, simdgroup s -> stream s / 2 and +// d0 + 8 (s % 2) .. + 8; LPR = 4 lanes per row (320 = 20 units of 16, 5 per lane). +// XP: [4 kinds][n_tok][n] then S16 [n_tok][n / 16] (units.metal's xprep layout). +template +kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* lo [[buffer(1)]], + device const float* xn [[buffer(2)]], device float* xp [[buffer(3)]], + constant GrArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float gate[8][4][16]; + threadgroup float msum[8][16]; + threadgroup uint notab[1]; + const int n = int(a.n_embd), c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3; + const int d = int(tg.x) * 16 + dl; + constexpr int UNIT = fmt::UNIT; + device const uchar* w = w_up + ulong(c * n + d) * (a.hc_lr / fmt::QK * fmt::BB); + const int U = int(a.hc_lr) / UNIT; + float acc[8]; + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (int u = li; u < U; u += 4) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, notab); + for (int t = 0; t < GR_NTOK; ++t) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + for (int k = 0; k < UNIT; k += 16) load16(lo + t * int(a.hc_lr) + UNIT * u + k, xs + k); + acc[t] += unit_combine(q, sc, xs, xsum); + } + } + for (int t = 0; t < GR_NTOK; ++t) { + const float v = lane_sum<4>(acc[t]); + if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // 16 threads per token: thread j -> d0 + j + const int t = tid / 16, j = tid & 15; + if (t < GR_NTOK) { + const int dd = int(tg.x) * 16 + j; + const int hc = int(a.hc); + float m = 0.0f; + for (int cc = 0; cc < hc; ++cc) m = fma(xn[(t * hc + cc) * n + dd], gate[t][cc][j], m); + m = m / float(hc); + const uint plane = a.n_tok * uint(n); + device float* raw = xp + t * n + dd; + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; + const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + raw[0] = m; + raw[plane] = m * q2s[j & 7]; + raw[2 * plane] = m * p4s[j & 3]; + raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; + msum[t][j] = m; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t < GR_NTOK && j == 0) { + // the per-16 sum in xprep's order (even and odd partial sums, then added) + float s0 = 0.0f, s1 = 0.0f; + for (int k = 0; k < 16; k += 2) { s0 += msum[t][k]; s1 += msum[t][k + 1]; } + xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + } +} + +#define GRD(S, K) template [[host_name("gr_down_" #S "_" #K)]] kernel decltype(gr_down_f) gr_down_f; \ + template [[host_name("gr_down_q8_" #S "_" #K)]] kernel decltype(gr_down_f) gr_down_f; \ + template [[host_name("gr_down_q8f_" #S "_" #K)]] kernel decltype(gr_down_f) gr_down_f; +GRD(8, 8) GRD(8, 4) GRD(4, 4) GRD(4, 2) GRD(16, 16) GRD(16, 8) +template [[host_name("gr_up")]] kernel decltype(gr_up_f) gr_up_f; +template [[host_name("gr_up_q8")]] kernel decltype(gr_up_f) gr_up_f; diff --git a/edge0-strata/metal/kernels/gr2.metal b/edge0-strata/metal/kernels/gr2.metal new file mode 100644 index 00000000..1253a11d --- /dev/null +++ b/edge0-strata/metal/kernels/gr2.metal @@ -0,0 +1,341 @@ +// metal/kernels/gr2.metal - the hyper-connection read in TWO dispatches (after quant/units/gr.metal), for BF16 HC +// weights (the exact path) or a Q8_0 copy of them (the STRATA_HC_Q8 option): +// +// gr_down3 grid (row blocks, hc streams). The RMSNorm is per stream, so a threadgroup owns one stream c: it forms +// R'_c = R_c + bo * 2 sigmoid(inj_prev_c / hc) (10 KB a token), its sum of squares and rs_c itself, +// stages xn_c = (R'_c * w_norm_c) * rs_c in threadgroup memory once, and every simdgroup dots its rows' +// stream-c quarter with it: D[t][row][c]. Row block 0 also writes R'_c to R_out (ping-pong, never in +// place) and rs_c. Rows 0..hc_lr-1 are w_down, hc_lr.. the inject rows. +// gr_up3 lo[k] = silu(sum_c D[t][k][c] / hc) formed in threadgroup memory, then as gr_up with xn = R' * w_norm +// * rs on the fly, writing the prepared activations. +// inj_prev of the NEXT read = sum_c D[t][hc_lr + c'][c] of this one (gr_down3 sums it, in c order). +// +// Per element the arithmetic is gr_norm's (fma(bo, g, R), then (R' * w) * rs); the K order differs (per stream, then +// the 4 streams in order). + +#ifndef P_TOKENT +#define P_TOKENT 532 // moe.metal's plan layout +#endif + +struct GrArgs3 { + uint n_embd; + uint hc; + uint hc_lr; + uint n_tok; // the prepared planes' row count (the plane stride) + uint apply; + float eps; + uint rb_down; // bytes per row of w_down / w_inject + uint rb_up; // bytes per row of w_up + uint comb; // 1: bo is not given; it is the MoE combine (moe_combine_sh's arithmetic) of buffers 10..13 + uint topk; +}; + +// bo[t][d], given or the MoE combine: sum_k ye[entry(t, k)][d] (k in order) + sigmoid(sg[t]) * ysh[t][d] +static inline float bo_at(device const float* bo, device const int* plan, device const float* ye, + device const float* ysh, device const float* sg, constant GrArgs3& a, int t, int d) { + const int n = int(a.n_embd); + if (!a.comb) return bo[t * n + d]; + float s = 0.0f; + for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * int(a.topk) + j] * n + d]; + const float gs = 1.0f / (1.0f + exp(-sg[t])); + return fma(gs, ysh[t * n + d], s); +} + +template +kernel void gr_down3(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* Rin [[buffer(2)]], device float* Rout [[buffer(3)]], + device const float* bo [[buffer(4)]], device const float* Dprev [[buffer(5)]], + device const float* wn [[buffer(6)]], device float* D [[buffer(7)]], + device float* rs_out [[buffer(8)]], constant GrArgs3& a [[buffer(9)]], + device const int* plan [[buffer(10)]], device const float* ye [[buffer(11)]], + device const float* ysh [[buffer(12)]], device const float* sgv [[buffer(13)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int UNIT = fmt::UNIT; + threadgroup float xs_tg[TT][2560]; + threadgroup float part[8][TT]; + threadgroup float g_tg[TT], rs_tg[TT]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 256); + const int n = int(a.n_embd), hc = int(a.hc), NIN = hc * n, c = int(tg.y); + const int nrows = int(a.hc_lr) + hc; + for (int t0 = 0; t0 < GR_NTOK; t0 += TT) { + const int T = min(TT, GR_NTOK - t0); + if (tid < T) { + float g = 0.0f; + if (a.apply) { + device const float* dp = Dprev + ((t0 + tid) * nrows + int(a.hc_lr) + c) * hc; + const float ip = ((dp[0] + dp[1]) + dp[2]) + dp[3]; + g = 2.0f / (1.0f + exp(-ip / float(hc))); + } + g_tg[tid] = g; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) { + device const float* rt = Rin + (t0 + t) * NIN + c * n; + float ss = 0.0f; + for (int d = tid; d < n; d += 256) { + float v = rt[d]; + if (a.apply) v = fma(bo_at(bo, plan, ye, ysh, sgv, a, t0 + t, d), g_tg[t], v); + xs_tg[t][d] = v; + ss = fma(v, v, ss); + if (tg.x == 0) Rout[(t0 + t) * NIN + c * n + d] = v; + } + ss = simd_sum(ss); + if (lane == 0) part[sg][t] = ss; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < T) { + float tot = 0.0f; + for (int k = 0; k < 8; ++k) tot += part[k][tid]; + const float rsv = rsqrt(tot / float(n) + a.eps); + rs_tg[tid] = rsv; + if (tg.x == 0) rs_out[(t0 + tid) * hc + c] = rsv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) + for (int d = tid; d < n; d += 256) xs_tg[t][d] = xs_tg[t][d] * wn[c * n + d] * rs_tg[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + // the rows: simdgroup sg holds RPS rows; lanes split the stream's units + int rows[RPS]; + device const uchar* wr[RPS]; + for (int k = 0; k < RPS; ++k) { + rows[k] = (int(tg.x) * 8 + sg) * RPS + k; + const int rr = min(rows[k], nrows - 1); + wr[k] = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; + } + float acc[RPS][TT]; + for (int k = 0; k < RPS; ++k) + for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; + const int U = n / UNIT, u0 = c * U; + for (int u = lane; u < U; u += 32) { + float q[RPS][UNIT]; + float4 sc[RPS]; + for (int k = 0; k < RPS; ++k) Dec::run(wr[k], u0 + u, q[k], sc[k], tab); + for (int t = 0; t < T; ++t) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + for (int e = 0; e < UNIT; ++e) xs[e] = xs_tg[t][u * UNIT + e]; + for (int k = 0; k < RPS; ++k) acc[k][t] += unit_combine(q[k], sc[k], xs, xsum); + } + } + for (int k = 0; k < RPS; ++k) + for (int t = 0; t < T; ++t) { + const float v = simd_sum(acc[k][t]); + if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +// grid (n_embd / 16), 256 threads. +template +kernel void gr_up3(device const uchar* w_up [[buffer(0)]], device const float* D [[buffer(1)]], + device const float* Rn [[buffer(2)]], device float* xp [[buffer(3)]], + constant GrArgs3& a [[buffer(4)]], device const float* wn [[buffer(5)]], + device const float* rs [[buffer(6)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT; + threadgroup float lo_tg[8][320]; + threadgroup float gate[8][4][16]; + threadgroup float msum[8][16]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 256); + const int n = int(a.n_embd), hc = int(a.hc), lr = int(a.hc_lr), nrows = lr + hc; + for (int e = tid; e < GR_NTOK * lr; e += 256) { + const int t = e / lr, k = e - t * lr; + device const float* dp = D + (t * nrows + k) * hc; + const float z = (((dp[0] + dp[1]) + dp[2]) + dp[3]) / float(hc); + lo_tg[t][k] = z / (1.0f + exp(-z)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3; + const int d = int(tg.x) * 16 + dl; + device const uchar* w = w_up + ulong(c * n + d) * a.rb_up; + const int U = lr / UNIT; + float acc[8]; + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (int u = li; u < U; u += 4) { + float q[UNIT]; + float4 sc; + Dec::run(w, u, q, sc, tab); + for (int t = 0; t < GR_NTOK; ++t) { + float xs[UNIT], xsum[2] = {0.0f, 0.0f}; + for (int e = 0; e < UNIT; ++e) xs[e] = lo_tg[t][UNIT * u + e]; + acc[t] += unit_combine(q, sc, xs, xsum); + } + } + for (int t = 0; t < GR_NTOK; ++t) { + const float v = lane_sum<4>(acc[t]); + if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int t = tid / 16, j = tid & 15; + if (t < GR_NTOK) { + const int dd = int(tg.x) * 16 + j; + float m = 0.0f; + for (int cc = 0; cc < hc; ++cc) { + const float xn = Rn[(t * hc + cc) * n + dd] * wn[cc * n + dd] * rs[t * hc + cc]; + m = fma(xn, gate[t][cc][j], m); + } + m = m / float(hc); + const uint plane = a.n_tok * uint(n); + device float* raw = xp + t * n + dd; + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, + 0.00006103515625f}; + const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + raw[0] = m; + raw[plane] = m * q2s[j & 7]; + raw[2 * plane] = m * p4s[j & 3]; + raw[3 * plane] = (j & 1) ? m * 0.00390625f : m; + msum[t][j] = m; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (t < GR_NTOK && j == 0) { + float s0 = 0.0f, s1 = 0.0f; + for (int k = 0; k < 16; k += 2) { s0 += msum[t][k]; s1 += msum[t][k + 1]; } + xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1; + } +} + +#define GD3(F, T, R) template [[host_name("gr_down3_" #F "_" #T "_" #R)]] kernel decltype(gr_down3) gr_down3; +GD3(30, 1, 1) GD3(30, 2, 1) GD3(30, 3, 1) GD3(30, 1, 2) GD3(30, 2, 2) GD3(30, 3, 2) +GD3(8, 1, 1) GD3(8, 2, 1) GD3(8, 3, 1) GD3(8, 1, 2) GD3(8, 2, 2) GD3(8, 3, 2) +template [[host_name("gr_up3_30")]] kernel decltype(gr_up3) gr_up3; +template [[host_name("gr_up3_8")]] kernel decltype(gr_up3) gr_up3; + +// BF16 only: as gr_down3, but each lane loads its rows' weights (5 units of 16 per row and lane: 2560 = 160 units) +// into registers before the norm phase, so the weight stream is in flight while the threadgroup forms xn. +template +kernel void gr_down3b(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]], + device const float* Rin [[buffer(2)]], device float* Rout [[buffer(3)]], + device const float* bo [[buffer(4)]], device const float* Dprev [[buffer(5)]], + device const float* wn [[buffer(6)]], device float* D [[buffer(7)]], + device float* rs_out [[buffer(8)]], constant GrArgs3& a [[buffer(9)]], + device const int* plan [[buffer(10)]], device const float* ye [[buffer(11)]], + device const float* ysh [[buffer(12)]], device const float* sgv [[buffer(13)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float xs_tg[TT][2560]; + threadgroup float part[8][TT]; + threadgroup float g_tg[TT], rs_tg[TT]; + const int n = int(a.n_embd), hc = int(a.hc), NIN = hc * n, c = int(tg.y); + const int nrows = int(a.hc_lr) + hc; + int rows[RPS]; + uint4 wr[RPS][5][2]; + for (int k = 0; k < RPS; ++k) { + rows[k] = (int(tg.x) * 8 + sg) * RPS + k; + const int rr = min(rows[k], nrows - 1); + device const uchar* w = rr < int(a.hc_lr) ? w_down + ulong(rr) * a.rb_down : w_inject + ulong(rr - int(a.hc_lr)) * a.rb_down; + device const uint4* p = (device const uint4*)(w + 2 * (c * n)); + for (int i = 0; i < 5; ++i) { + wr[k][i][0] = p[2 * (lane + 32 * i)]; + wr[k][i][1] = p[2 * (lane + 32 * i) + 1]; + } + } + for (int t0 = 0; t0 < GR_NTOK; t0 += TT) { + const int T = min(TT, GR_NTOK - t0); + if (tid < T) { + float g = 0.0f; + if (a.apply) { + device const float* dp = Dprev + ((t0 + tid) * nrows + int(a.hc_lr) + c) * hc; + const float ip = ((dp[0] + dp[1]) + dp[2]) + dp[3]; + g = 2.0f / (1.0f + exp(-ip / float(hc))); + } + g_tg[tid] = g; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) { + device const float* rt = Rin + (t0 + t) * NIN + c * n; + float ss = 0.0f; + for (int d = tid; d < n; d += 256) { + float v = rt[d]; + if (a.apply) v = fma(bo_at(bo, plan, ye, ysh, sgv, a, t0 + t, d), g_tg[t], v); + xs_tg[t][d] = v; + ss = fma(v, v, ss); + if (tg.x == 0) Rout[(t0 + t) * NIN + c * n + d] = v; + } + ss = simd_sum(ss); + if (lane == 0) part[sg][t] = ss; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < T) { + float tot = 0.0f; + for (int k = 0; k < 8; ++k) tot += part[k][tid]; + const float rsv = rsqrt(tot / float(n) + a.eps); + rs_tg[tid] = rsv; + if (tg.x == 0) rs_out[(t0 + tid) * hc + c] = rsv; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int t = 0; t < T; ++t) + for (int d = tid; d < n; d += 256) xs_tg[t][d] = xs_tg[t][d] * wn[c * n + d] * rs_tg[t]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[RPS][TT]; + for (int k = 0; k < RPS; ++k) + for (int t = 0; t < TT; ++t) acc[k][t] = 0.0f; + for (int i = 0; i < 5; ++i) { + const int u = lane + 32 * i; + for (int k = 0; k < RPS; ++k) { + const uint v[8] = {wr[k][i][0].x, wr[k][i][0].y, wr[k][i][0].z, wr[k][i][0].w, + wr[k][i][1].x, wr[k][i][1].y, wr[k][i][1].z, wr[k][i][1].w}; + float q[16]; + for (int e = 0; e < 8; ++e) { + q[2 * e] = as_type(v[e] << 16); + q[2 * e + 1] = as_type(v[e] & 0xFFFF0000u); + } + const float4 sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + for (int t = 0; t < T; ++t) { + float xs[16], xsum[2] = {0.0f, 0.0f}; + for (int e = 0; e < 16; ++e) xs[e] = xs_tg[t][u * 16 + e]; + acc[k][t] += unit_combine<16>(q, sc, xs, xsum); + } + } + } + for (int k = 0; k < RPS; ++k) + for (int t = 0; t < T; ++t) { + const float v = simd_sum(acc[k][t]); + if (lane == 0 && rows[k] < nrows) D[((t0 + t) * nrows + rows[k]) * hc + c] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} +#define GD3B(T, R) template [[host_name("gr_down3b_" #T "_" #R)]] kernel decltype(gr_down3b) gr_down3b; +GD3B(1, 1) GD3B(2, 1) GD3B(3, 1) GD3B(1, 2) GD3B(2, 2) GD3B(3, 2) + +// gr_norm_comb: gr_norm (gr.metal) with the sublayer output bo formed on the fly as the MoE combine - bo[t][d] = +// sum_k ye[entry(t, k)][d] (k in order) + sigmoid(sg[t]) * ysh[t][d], moe_combine_sh's arithmetic - so the combine's +// dispatch goes away; R', xn bitwise gr_norm's after moe_combine_sh. One threadgroup per (stream c, token t). +kernel void gr_norm_comb(device float* R [[buffer(0)]], device const float* ye [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + device const int* plan [[buffer(6)]], device const float* ysh [[buffer(7)]], + device const float* sgv [[buffer(8)]], constant uint& topk [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + const float gs = 1.0f / (1.0f + exp(-sgv[t])); + float ss = 0.0f; + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + float s = 0.0f; + for (uint j = 0; j < topk; ++j) s += ye[plan[P_TOKENT + t * topk + j] * int(n) + d]; + const float bo = fma(gs, ysh[t * n + d], s); + v = fma(bo, g, v); + r[d] = v; + } + ss = fma(v, v, ss); + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device float* o = xn + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + for (uint d = tid; d < n; d += 256) o[d] = r[d] * w[d] * rs; +} diff --git a/edge0-strata/metal/kernels/gridbar.metal b/edge0-strata/metal/kernels/gridbar.metal new file mode 100644 index 00000000..5cd9a0fa --- /dev/null +++ b/edge0-strata/metal/kernels/gridbar.metal @@ -0,0 +1,54 @@ +// metal/kernels/gridbar.metal - feasibility probe for a persistent kernel: phases separated by an in-kernel grid +// barrier (a device atomic counter) instead of dispatch boundaries. Metal does not promise that all threadgroups of +// a grid are resident at once; the probe launches at most one threadgroup per GPU core. +// +// Each phase p: every thread streams its share of chunk p (n uint4), and threadgroup g writes a token value +// (p * 1000 + g) to flag[p][g]; after the barrier, phase p + 1 checks every flag of phase p (written by the OTHER +// threadgroups) and counts mismatches into bad[0]. + +static inline void grid_barrier(device atomic_uint* bar, uint target, ushort tid) { + threadgroup_barrier(mem_flags::mem_device); +#ifdef GB_FENCE + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); +#endif + if (tid == 0) { + atomic_fetch_add_explicit(bar, 1u, memory_order_relaxed); + while (atomic_load_explicit(bar, memory_order_relaxed) < target) {} + } + threadgroup_barrier(mem_flags::mem_device); +#ifdef GB_FENCE + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); +#endif +} + +kernel void k_phases(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + device atomic_uint* bar [[buffer(2)]], device volatile uint* flag [[buffer(3)]], + device atomic_uint* bad [[buffer(4)]], constant uint& n [[buffer(5)]], + constant uint& phases [[buffer(6)]], + uint tgpx [[threadgroup_position_in_grid]], uint ntgx [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], uint gid [[thread_position_in_grid]], + uint nt [[threads_per_grid]]) { + const uint g = tgpx, ntg = ntgx; + uint4 acc = 0; + for (uint p = 0; p < phases; ++p) { + if (p > 0 && tid < ntg) { + const uint v = flag[(p - 1) * ntg + tid]; + if (v != (p - 1) * 1000 + tid) atomic_fetch_add_explicit(bad, 1u, memory_order_relaxed); + } + device const uint4* s = src + ulong(p) * n; + for (uint i = gid; i < n; i += nt) acc += s[i]; + if (tid == 0) flag[p * ntg + g] = p * 1000 + g; + grid_barrier(bar, (p + 1) * ntg, tid); + } + out[gid] = acc.x ^ acc.y ^ acc.z ^ acc.w; +} + +// one phase per dispatch, for comparison +kernel void k_phase1(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& p [[buffer(3)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 acc = 0; + device const uint4* s = src + ulong(p) * n; + for (uint i = gid; i < n; i += nt) acc += s[i]; + out[gid] = acc.x ^ acc.y ^ acc.z ^ acc.w; +} diff --git a/edge0-strata/metal/kernels/latent_prompt.metal b/edge0-strata/metal/kernels/latent_prompt.metal new file mode 100644 index 00000000..58185edf --- /dev/null +++ b/edge0-strata/metal/kernels/latent_prompt.metal @@ -0,0 +1,54 @@ +// D/U scratch rows are group-major. Expert rows use the prompt planner's +// padded order; each tile identifies one expert and each valid row one token. +struct LatentPromptArgs { + uint tokens, groups, hidden, width, topk, tile_rows; +}; + +kernel void latent_prompt_gather(device const float* d [[buffer(0)]], + device const float* inverse [[buffer(1)]], + device const int* labels [[buffer(2)]], + device const int2* tiles [[buffer(3)]], + device const int* tokens [[buffer(4)]], + device half* x [[buffer(5)]], + constant LatentPromptArgs& a [[buffer(6)]], + uint row [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + const int token = tokens[row]; + device half* out = x + ulong(row) * a.width; + if (token < 0) { + for (uint k = tid; k < a.width; k += 256) out[k] = half(0.0f); + return; + } + const uint group = labels[tiles[row / a.tile_rows].x]; + device const float* in = d + (ulong(group) * a.tokens + uint(token)) * a.width; + for (uint k = tid; k < a.width; k += 256) out[k] = half(in[k] * inverse[k]); +} + +kernel void latent_prompt_pool(device const float* down [[buffer(0)]], + device const int* positions [[buffer(1)]], + device const int* ids [[buffer(2)]], + device const float* weights [[buffer(3)]], + device const int* labels [[buffer(4)]], + device half* pooled [[buffer(5)]], + constant LatentPromptArgs& a [[buffer(6)]], + uint i [[thread_position_in_grid]]) { + const uint row = i / a.width, k = i % a.width; + if (row >= a.groups * a.tokens) return; + const uint group = row / a.tokens, entry = (row % a.tokens) * a.topk; + float sum = 0.0f; + for (uint j = entry; j < entry + a.topk; ++j) + if (uint(labels[ids[j]]) == group) + sum = fma(weights[j], down[ulong(positions[j]) * a.width + k], sum); + pooled[i] = half(sum); +} + +kernel void latent_prompt_sum(device const float* projected [[buffer(0)]], + device float* routed [[buffer(1)]], + constant LatentPromptArgs& a [[buffer(2)]], + uint i [[thread_position_in_grid]]) { + const uint stride = a.tokens * a.hidden; + if (i >= stride) return; + float sum = 0.0f; + for (uint group = 0; group < a.groups; ++group) sum += projected[group * stride + i]; + routed[i] = sum; +} diff --git a/edge0-strata/metal/kernels/latent_prompt_q3_stage.metal b/edge0-strata/metal/kernels/latent_prompt_q3_stage.metal new file mode 100644 index 00000000..39899231 --- /dev/null +++ b/edge0-strata/metal/kernels/latent_prompt_q3_stage.metal @@ -0,0 +1,71 @@ +// Cache the prefill GEMM's existing FP16 weight operand once per layer. +// Q3_K bytes, scales, dequantization and the K tile of the GEMM stay the same. +struct LatentPromptQ3Args { uint width, ff, row_bytes, expert_bytes; }; + +kernel void latent_prompt_q3_stage(device const uchar* experts [[buffer(0)]], + device half* weights [[buffer(1)]], + constant LatentPromptQ3Args& a [[buffer(2)]], + uint2 gid [[thread_position_in_grid]]) { + const uint units = a.width / 32; + if (gid.x >= 2 * a.ff * units) return; + const uint row = gid.x / units, unit = gid.x % units; + const uint expert = gid.y; + const device uchar* source = experts + ulong(expert) * a.expert_bytes + row * a.row_bytes; + device half* output = weights + ulong(expert * 2 * a.ff + row) * a.width; + threadgroup uint table[1]; + float values[32]; + int lo, hi; + DW::run(source, unit, 0, values, lo, hi, table); + for (uint j = 0; j < 16; ++j) { + output[lo + j] = half(values[j]); + output[hi + j] = half(values[16 + j]); + } +} + +kernel void latent_prompt_gu_f16(device const half* weights [[buffer(0)]], + device const half* input [[buffer(1)]], + device half* middle [[buffer(2)]], + device const int2* tiles [[buffer(3)]], + constant MmArgs& a [[buffer(4)]], + uint2 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + constexpr int rows = 32, cols = 64, tile_k = 32; + threadgroup float output[rows * cols]; + const int2 tile = tiles[tg.y]; + const int row0 = tile.y, f0 = int(tg.x) * cols; + device half* gate = (device half*)weights + ulong(tile.x * 2 * a.n_ff) * a.n_embd; + tensor, tensor_inline> x( + (device half*)input, dextents(a.n_embd, 1 << 30)); + tensor, tensor_inline> g(gate, dextents(a.n_embd, a.n_ff)); + tensor, tensor_inline> u( + gate + ulong(a.n_ff) * a.n_embd, dextents(a.n_embd, a.n_ff)); + constexpr auto desc = matmul2d_descriptor(rows, cols, tile_k, false, true, false, + matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> op; + auto xa = x.template slice(0, row0); + auto wg = g.template slice(0, f0); + auto sum_g = op.template get_destination_cooperative_tensor(); + auto sum_u = op.template get_destination_cooperative_tensor(); + for (ushort i = 0; i < sum_g.get_capacity(); ++i) { + if (!sum_g.is_valid_element(i)) continue; + sum_g[i] = 0.0f; + sum_u[i] = 0.0f; + } + for (int k = 0; k < int(a.n_embd); k += tile_k) { + xa = x.template slice(k, row0); + wg = g.template slice(k, f0); + auto wu = u.template slice(k, f0); + op.run(xa, wg, sum_g); + op.run(xa, wu, sum_u); + } + for (ushort i = 0; i < sum_g.get_capacity(); ++i) { + if (!sum_g.is_valid_element(i)) continue; + const float value = sum_g[i]; + sum_g[i] = (value / (1.0f + exp(-value))) * sum_u[i]; + } + tensor, tensor_inline> y(output, dextents(cols, rows)); + sum_g.store(y); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int i = tid; i < rows * cols; i += 64) + middle[ulong(row0 + i / cols) * a.n_ff + f0 + i % cols] = half(output[i]); +} diff --git a/edge0-strata/metal/kernels/latent_routed.metal b/edge0-strata/metal/kernels/latent_routed.metal new file mode 100644 index 00000000..f8245db2 --- /dev/null +++ b/edge0-strata/metal/kernels/latent_routed.metal @@ -0,0 +1,181 @@ +// The latent routed FFN. Reuse the quantized-row decoders from quant/units.metal. +struct LatentRoutedArgs { + uint n_tok, groups, hidden, latent, ff, topk, plan_token_entry; + uint u_row, gu_row, down_row, up_off, down_off, expert_bytes; +}; + +kernel void latent_colscale(device float* projected [[buffer(0)]], + device const float* inverse [[buffer(1)]], + constant LatentRoutedArgs& a [[buffer(2)]], + uint i [[thread_position_in_grid]]) { + if (i < a.n_tok * a.groups * a.latent) projected[i] *= inverse[i % a.latent]; +} + +template +static inline float latent_q3k_dot(device const uchar* row, device const float* x, + uint width, ushort lane, threadgroup const uint* table) { + float sum = 0.0f; + for (uint unit = lane; unit < width / 32; unit += Lanes) { + float codes[32]; + float4 scale; + Dec::run(row, unit, codes, scale, table); + float low = 0.0f, high = 0.0f; + for (uint j = 0; j < 16; ++j) { + // Dec leaves odd elements in the high byte of each packed ushort. + const float unpack = (j & 1) ? 1.0f / 256.0f : 1.0f; + low = fma(codes[j] * unpack - 4.0f, x[32 * unit + j], low); + high = fma(codes[j + 16] * unpack - 4.0f, x[32 * unit + j + 16], high); + } + sum = fma(scale.x, low, fma(scale.y, high, sum)); + } + return lane_sum(sum); +} + +template +static inline float latent_32_dot(device const uchar* row, device const float* x, + uint width, ushort lane, threadgroup const uint* table) { + float sum = 0.0f; + for (uint unit = lane; unit < width / 32; unit += Lanes) { + float codes[32]; + const float scale = deq(row + unit * fmt::BB, 0, codes, table); + float dot = 0.0f; + for (uint j = 0; j < 32; ++j) dot = fma(codes[j], x[32 * unit + j], dot); + sum = fma(scale, dot, sum); + } + return lane_sum(sum); +} + +template +kernel void latent_gate_up(device const uchar* experts [[buffer(0)]], + device const int* labels [[buffer(1)]], + device const int* ids [[buffer(2)]], + device const float* projected [[buffer(3)]], + device float* middle [[buffer(4)]], + constant LatentRoutedArgs& a [[buffer(5)]], + uint2 tg [[threadgroup_position_in_grid]], + ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup uint table[1]; // Q3_K's decoder does not use a lookup table + const uint entry = tg.y, row = (4 * tg.x + sg) * (32 / Lanes) + lane / Lanes; + if (row >= a.ff) return; + const uint expert = ids[entry], group = labels[expert]; + const uint token = entry / a.topk; + device const float* x = projected + (token * a.groups + group) * a.latent; + device const uchar* gate = experts + ulong(expert) * a.expert_bytes + row * a.gu_row; + const float g = latent_q3k_dot(gate, x, a.latent, lane % Lanes, table); + const float u = latent_q3k_dot(gate + a.up_off, x, a.latent, lane % Lanes, table); + if (lane % Lanes == 0) middle[entry * a.ff + row] = (g / (1.0f + exp(-g))) * u; +} + +template +kernel void latent_down(device const uchar* experts [[buffer(0)]], + device const int* ids [[buffer(1)]], + device const float* middle [[buffer(2)]], + device float* output [[buffer(3)]], + constant LatentRoutedArgs& a [[buffer(4)]], + uint2 tg [[threadgroup_position_in_grid]], + ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint table[fmt::TAB]; + load_tab(table, tid, 128); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint entry = tg.y, row = (4 * tg.x + sg) * (32 / Lanes) + lane / Lanes; + if (row >= a.latent) return; + device const uchar* w = experts + ulong(ids[entry]) * a.expert_bytes + a.down_off + row * a.down_row; + const float v = latent_32_dot(w, middle + entry * a.ff, a.ff, lane % Lanes, table); + if (lane % Lanes == 0) output[entry * a.latent + row] = v; +} + +kernel void latent_u(device const uchar* weights [[buffer(0)]], + device const int* labels [[buffer(1)]], + device const int* ids [[buffer(2)]], + device const float* route_weight [[buffer(3)]], + device const int* plan [[buffer(4)]], + device const float* latent_rows [[buffer(5)]], + device float* output [[buffer(6)]], + constant LatentRoutedArgs& a [[buffer(7)]], + uint2 tg [[threadgroup_position_in_grid]], + ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup uint table[1]; // Q4_0's decoder does not use a lookup table + const uint entry = tg.y, row = 4 * tg.x + sg; + if (row >= a.hidden) return; + const uint group = labels[ids[entry]]; + const uint slot = plan[a.plan_token_entry + entry]; + device const uchar* w = weights + (group * a.hidden + row) * a.u_row; + const float v = latent_32_dot(w, latent_rows + entry * a.latent, a.latent, lane, table); + if (lane == 0) output[slot * a.hidden + row] = v * route_weight[entry]; +} + +// U is linear. Pool route-weighted expert outputs within a group before +// applying that group's U, then sum the projected groups with the shared expert. +kernel void latent_group_reduce(device const int* labels [[buffer(0)]], + device const int* ids [[buffer(1)]], + device const float* route_weight [[buffer(2)]], + device const float* latent_rows [[buffer(3)]], + device float* group_rows [[buffer(4)]], + device uint* active [[buffer(5)]], + constant LatentRoutedArgs& a [[buffer(6)]], + uint i [[thread_position_in_grid]]) { + const uint slot = i / a.latent; + if (slot >= a.n_tok * a.groups) return; + const uint column = i % a.latent, token = slot / a.groups, group = slot % a.groups; + float sum = 0.0f; + bool used = false; + for (uint rank = 0; rank < a.topk; ++rank) { + const uint entry = token * a.topk + rank; + if (uint(labels[ids[entry]]) == group) { + sum = fma(route_weight[entry], latent_rows[entry * a.latent + column], sum); + used = true; + } + } + group_rows[i] = sum; + if (column == 0) active[slot] = used; +} + +template +kernel void latent_group_u(device const uchar* weights [[buffer(0)]], + device const float* group_rows [[buffer(1)]], + device const uint* active [[buffer(2)]], + device float* output [[buffer(3)]], + constant LatentRoutedArgs& a [[buffer(4)]], + uint2 tg [[threadgroup_position_in_grid]], + ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup uint table[1]; + const uint slot = tg.y, row = (4 * tg.x + sg) * (32 / Lanes) + lane / Lanes; + if (row >= a.hidden) return; + float y = 0.0f; + if (active[slot]) { + const uint group = slot % a.groups; + device const uchar* w = weights + (group * a.hidden + row) * a.u_row; + y = latent_32_dot(w, group_rows + slot * a.latent, a.latent, lane % Lanes, table); + } + if (lane % Lanes == 0) output[slot * a.hidden + row] = y; +} + +template [[host_name("latent_gate_up")]] kernel decltype(latent_gate_up<32>) latent_gate_up<32>; +template [[host_name("latent_down")]] kernel decltype(latent_down<32>) latent_down<32>; +template [[host_name("latent_group_u")]] kernel decltype(latent_group_u<32>) latent_group_u<32>; +// Split SIMD-groups into independent rows to avoid inactive lanes in narrow dots. +template [[host_name("latent_gate_up_8")]] kernel decltype(latent_gate_up<8>) latent_gate_up<8>; +template [[host_name("latent_gate_up_4")]] kernel decltype(latent_gate_up<4>) latent_gate_up<4>; +template [[host_name("latent_down_4")]] kernel decltype(latent_down<4>) latent_down<4>; +template [[host_name("latent_group_u_8")]] kernel decltype(latent_group_u<8>) latent_group_u<8>; +template [[host_name("latent_group_u_4")]] kernel decltype(latent_group_u<4>) latent_group_u<4>; + +kernel void latent_group_combine(device const float* grouped [[buffer(0)]], + device const float* shared [[buffer(1)]], + device const float* shared_gate [[buffer(2)]], + device float* output [[buffer(3)]], + constant LatentRoutedArgs& a [[buffer(4)]], + uint i [[thread_position_in_grid]]) { + const uint token = i / a.hidden, row = i % a.hidden; + if (token >= a.n_tok) return; + float sum = 0.0f; + for (uint group = 0; group < a.groups; ++group) + sum += grouped[(token * a.groups + group) * a.hidden + row]; + const float g = 1.0f / (1.0f + exp(-shared_gate[token])); + output[i] = fma(g, shared[i], sum); +} diff --git a/edge0-strata/metal/kernels/mm_bench.metal b/edge0-strata/metal/kernels/mm_bench.metal new file mode 100644 index 00000000..c04659fc --- /dev/null +++ b/edge0-strata/metal/kernels/mm_bench.metal @@ -0,0 +1,85 @@ +// metal/kernels/mm_bench.metal - matrix throughput probes (mm_bench.cpp): C[M][N] = A[M][K] . B[N][K]^T, fp32 out. +// mpp_mm Metal 4 MetalPerformancePrimitives matmul2d (the M5's neural accelerators), a TM x TN +// tile of C per threadgroup, K dynamic, A and B straight from device memory +// sg_mm simdgroup_matrix 8x8, a 32 x 32 tile per simdgroup (4 x 4 accumulators), 4 simdgroups +#include +using namespace mpp::tensor_ops; + +struct MmArgs { int M, N, K; }; + +template +kernel void mpp_mm(device T* A [[buffer(0)]], device T* B [[buffer(1)]], device float* C [[buffer(2)]], + constant MmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) { + // extents are (inner, outer): A is M rows of K, B is N rows of K, C is M rows of N + tensor, tensor_inline> tA(A, dextents(a.K, a.M)); + tensor, tensor_inline> tB(B, dextents(a.K, a.N)); + tensor, tensor_inline> tC(C, dextents(a.N, a.M)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mB = tB.slice(0, int(tg.x) * TN); + auto mC = tC.slice(int(tg.x) * TN, int(tg.y) * TM); + op.run(mA, mB, mC); +} + +#define MPP(T, M, N, S) template [[host_name("mpp_mm_" #T "_" #M "_" #N "_" #S)]] kernel decltype(mpp_mm) mpp_mm; +MPP(half, 64, 32, 4) MPP(half, 64, 64, 4) MPP(half, 128, 64, 4) MPP(half, 32, 64, 4) MPP(half, 64, 128, 4) +MPP(half, 16, 64, 2) MPP(half, 128, 128, 4) +MPP(half, 64, 16, 2) MPP(half, 64, 16, 4) MPP(half, 128, 16, 4) MPP(half, 128, 32, 4) MPP(half, 64, 32, 2) MPP(half, 32, 16, 2) +MPP(half, 32, 32, 2) MPP(half, 16, 32, 1) MPP(half, 16, 16, 1) MPP(half, 128, 48, 4) MPP(half, 64, 48, 4) +MPP(float, 64, 32, 4) MPP(float, 64, 64, 4) MPP(float, 128, 64, 4) +MPP(bfloat, 64, 64, 4) MPP(bfloat, 128, 64, 4) + +// simdgroup_matrix: threadgroup tile 64 x 64 (4 simdgroups x 32 x 32), K in steps of 8, operands from device memory +template +kernel void sg_mm(device const T* A [[buffer(0)]], device const T* B [[buffer(1)]], device float* C [[buffer(2)]], + constant MmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + const int m0 = int(tg.y) * 64 + (sg >> 1) * 32, n0 = int(tg.x) * 64 + (sg & 1) * 32; + simdgroup_matrix acc[4][4]; + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) acc[i][j] = make_filled_simdgroup_matrix(0.0f); + for (int k = 0; k < a.K; k += 8) { + simdgroup_matrix ma[4], mb[4]; + for (int i = 0; i < 4; ++i) simdgroup_load(ma[i], A + (m0 + 8 * i) * a.K + k, a.K); + for (int j = 0; j < 4; ++j) simdgroup_load(mb[j], B + (n0 + 8 * j) * a.K + k, a.K, ulong2(0, 0), true); + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) simdgroup_multiply_accumulate(acc[i][j], ma[i], mb[j], acc[i][j]); + } + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) simdgroup_store(acc[i][j], C + (m0 + 8 * i) * a.N + n0 + 8 * j, a.N); +} +template [[host_name("sg_mm_half")]] kernel decltype(sg_mm) sg_mm; +template [[host_name("sg_mm_float")]] kernel decltype(sg_mm) sg_mm; + + +// the element layout of single-simdgroup matmul2d cooperative tensors (right input, left input, destination) for a +// TM x TN x TK descriptor: out[((which * 32 + lane) * 128 + i) * 4 + (0, 1)] = the two indices, [2] = capacity +template +kernel void coop_layout(device int* out [[buffer(0)]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + auto r = op.template get_right_input_cooperative_tensor(); + auto l = op.template get_left_input_cooperative_tensor(); + auto dd = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < dd.get_capacity() && i < 128; ++i) { + auto ix = dd.get_multidimensional_index(i); + device int* o = out + ((2 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = dd.get_capacity(); + } + for (uint16_t i = 0; i < r.get_capacity() && i < 128; ++i) { + auto ix = r.get_multidimensional_index(i); + device int* o = out + ((0 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = r.get_capacity(); + } + for (uint16_t i = 0; i < l.get_capacity() && i < 128; ++i) { + auto ix = l.get_multidimensional_index(i); + device int* o = out + ((1 * 32 + lane) * 128 + i) * 4; + o[0] = ix[0]; o[1] = ix[1]; o[2] = l.get_capacity(); + } +} +template [[host_name("coop_layout_16_32_16")]] kernel decltype(coop_layout<16, 32, 16>) coop_layout<16, 32, 16>; +template [[host_name("coop_layout_32_32_16")]] kernel decltype(coop_layout<32, 32, 16>) coop_layout<32, 32, 16>; +template [[host_name("coop_layout_16_64_32")]] kernel decltype(coop_layout<16, 64, 32>) coop_layout<16, 64, 32>; +template [[host_name("coop_layout_16_32_32")]] kernel decltype(coop_layout<16, 32, 32>) coop_layout<16, 32, 32>; diff --git a/edge0-strata/metal/kernels/mma.metal b/edge0-strata/metal/kernels/mma.metal new file mode 100644 index 00000000..9bbb27e2 --- /dev/null +++ b/edge0-strata/metal/kernels/mma.metal @@ -0,0 +1,216 @@ +// metal/kernels/mma.metal - a GEMV for windows of tokens on the simdgroup matrix units (after quant/units.metal). +// +// y[t][r] = W[r] . x[t] for NTOK (1..8) tokens: each simdgroup owns 8 rows; the K axis runs in chunks of 128. Per +// chunk the simdgroup decodes its 8 x 128 weights - with ref_quant's arithmetic, so a tile element is exactly the +// dequantized weight - into threadgroup memory, then 16 simdgroup_multiply_accumulate steps take A = the 8 x 8 weight +// tile and B = the 8 (k) x 8 (token) slice of x, loaded transposed straight from device memory (x has 8 rows: the +// tokens past the window are padding, and an output column depends only on its own token). +// +// Activations are the raw plane (no pre-scaling); n_in must be a multiple of 128. Every format's units sit inside +// one 128-wide chunk: Q4_K/Q5_K units span 64 positions, Q6_K units 128, the others are contiguous. + +// ---- decoders to actual weights: w[UNIT] and the chunk offsets of its two halves (p0: w[0..H), p1: w[H..UNIT)) ---- +template struct DW { + // generic: w = s * q (quant.metal's deq is ref_quant's order: the scale times the code), contiguous + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT, UPB = fmt::QK / UNIT; + float q[UNIT]; + const float s = deq(row + (u / UPB) * fmt::BB, u % UPB, q, tab); + for (int j = 0; j < UNIT; ++j) w[j] = q[j] * s; + p0 = u * UNIT - k0; + p1 = p0 + UNIT / 2; + } +}; +// Q2_0: (code - 1) * d +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + const float d = rdh(b); + for (int j = 0; j < 8; ++j) { + w[j] = float(int((a >> (2 * j)) & 3) - 1) * d; + w[8 + j] = float(int((c >> (2 * j)) & 3) - 1) * d; + } + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint* tab) { + device const uchar* b = row + u * 24; + float q[32]; + const float d = deq(b, 0, q, tab), m = rdh(b + 2); + for (int j = 0; j < 32; ++j) w[j] = q[j] * d + m; + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const float* p = (device const float*)(row + 64 * u); + for (int j = 0; j < 16; ++j) w[j] = p[j]; + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uint4* p = (device const uint4*)(row + 32 * u); + const uint4 a = p[0], c = p[1]; + const uint v[8] = {a.x, a.y, a.z, a.w, c.x, c.y, c.z, c.w}; + for (int k = 0; k < 8; ++k) { + w[2 * k] = as_type(v[k] << 16); + w[2 * k + 1] = as_type(v[k] & 0xFFFF0000u); + } + p0 = 16 * u - k0; + p1 = p0 + 8; + } +}; +// Q4_K / Q5_K: d1 * code - m1 (ref: d * sc, dmin * m), runs at 256b + 64c + 16m and + 32 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 144; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + const float d1 = d * s1, mm1 = dm * m1, d2 = d * s2, mm2 = dm * m2; + device const uchar* q = b + 16 + 32 * c + 16 * m; + for (int j = 0; j < 16; ++j) { + w[j] = d1 * float(q[j] & 15) - mm1; + w[16 + j] = d2 * float(q[j] >> 4) - mm2; + } + p0 = 256 * (u >> 3) + 64 * c + 16 * m - k0; + p1 = p0 + 32; + } +}; +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 176; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + const float d1 = d * s1, mm1 = dm * m1, d2 = d * s2, mm2 = dm * m2; + device const uchar* ql = b + 48 + 32 * c + 16 * m; + device const uchar* qh = b + 16 + 16 * m; + const uchar u1 = uchar(1 << (2 * c)), u2 = uchar(2 << (2 * c)); + for (int j = 0; j < 16; ++j) { + w[j] = d1 * float((ql[j] & 15) + ((qh[j] & u1) ? 16 : 0)) - mm1; + w[16 + j] = d2 * float((ql[j] >> 4) + ((qh[j] & u2) ? 16 : 0)) - mm2; + } + p0 = 256 * (u >> 3) + 64 * c + 16 * m - k0; + p1 = p0 + 32; + } +}; +// Q6_K: d * sc * (code - 32), runs at 256b + 128n + 32p + 16m and + 64 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 210; + const int n = (u >> 2) & 1, m = (u >> 1) & 1, p = u & 1; + device const uchar* ql = b + 64 * n + 32 * p + 16 * m; + device const uchar* qh = b + 128 + 32 * n + 16 * m; + device const char* sc = (device const char*)(b + 192); + const float d = rdh(b + 208); + const float ds1 = d * float(sc[8 * n + m + 2 * p]), ds2 = d * float(sc[8 * n + m + 2 * p + 4]); + for (int j = 0; j < 16; ++j) { + const int c1 = int((ql[j] & 15) | (((qh[j] >> (2 * p)) & 3) << 4)) - 32; + const int c2 = int((ql[j] >> 4) | (((qh[j] >> (2 * p + 4)) & 3) << 4)) - 32; + w[j] = ds1 * float(c1); + w[16 + j] = ds2 * float(c2); + } + p0 = 256 * (u >> 3) + 128 * n + 32 * p + 16 * m - k0; + p1 = p0 + 64; + } +}; +// Q3_K: d_all * (sc - 32) * (code - (hbit ? 0 : 4)), contiguous 32 +template <> struct DW { + static inline void run(device const uchar* row, int u, int k0, thread float* w, thread int& p0, thread int& p1, + threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 110; + const int n = (u >> 2) & 1, j = u & 3; + device const uchar* q = b + 32 + 32 * n; + device const uchar* hm = b; + device const uchar* s = b + 96; + const int i1 = 8 * n + 2 * j, i2 = i1 + 1; + const int l1 = (i1 < 8 ? s[i1] & 0xF : s[i1 - 8] >> 4) | (((s[8 + (i1 & 3)] >> (2 * (i1 >> 2))) & 3) << 4); + const int l2 = (i2 < 8 ? s[i2] & 0xF : s[i2 - 8] >> 4) | (((s[8 + (i2 & 3)] >> (2 * (i2 >> 2))) & 3) << 4); + const float d = rdh(b + 108); + const float dl1 = d * float(l1 - 32), dl2 = d * float(l2 - 32); + const uchar hb = uchar(1 << (4 * n + j)); + for (int l = 0; l < 16; ++l) { + w[l] = dl1 * float(int((q[l] >> (2 * j)) & 3) - ((hm[l] & hb) ? 0 : 4)); + w[16 + l] = dl2 * float(int((q[l + 16] >> (2 * j)) & 3) - ((hm[l + 16] & hb) ? 0 : 4)); + } + p0 = 32 * u - k0; + p1 = p0 + 16; + } +}; + +struct MmaArgs { + uint n_in; + uint n_out; + uint row_bytes; + uint x_stride; // floats between the 8 token rows of x + uint y_stride; +}; +constant int MMA_NTOK [[function_constant(4)]]; + +template +kernel void gemv_mma(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* y [[buffer(2)]], constant MmaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, UPC = 128 / UNIT, H = UNIT / 2; + threadgroup float wt[SG][8 * 128]; + threadgroup float ct[SG][64]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int row0 = (int(tg.x) * SG + sg) * 8; + simdgroup_float8x8 C = make_filled_simdgroup_matrix(0.0f); + threadgroup float* w8 = wt[sg]; + for (int k0 = 0; k0 < int(a.n_in); k0 += 128) { + for (int it = lane; it < 8 * UPC; it += 32) { + const int r = it / UPC, uu = it - r * UPC; + const int rr = min(row0 + r, int(a.n_out) - 1); + float w[UNIT]; + int p0, p1; + DW::run(W + ulong(rr) * a.row_bytes, k0 / UNIT + uu, k0, w, p0, p1, tab); + for (int j = 0; j < H; ++j) { + w8[r * 128 + p0 + j] = w[j]; + w8[r * 128 + p1 + j] = w[H + j]; + } + } + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int k8 = 0; k8 < 128; k8 += 8) { + simdgroup_float8x8 A, B; + simdgroup_load(A, w8 + k8, 128); + simdgroup_load(B, x + k0 + k8, a.x_stride, ulong2(0, 0), true); + simdgroup_multiply_accumulate(C, A, B, C); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + simdgroup_store(C, ct[sg], 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int e = lane; e < 64; e += 32) { + const int r = e >> 3, t = e & 7; + if (t < MMA_NTOK && row0 + r < int(a.n_out)) y[t * int(a.y_stride) + row0 + r] = ct[sg][e]; + } +} + +#define MM(F, S) template [[host_name("gemv_mma_" #F "_" #S)]] kernel decltype(gemv_mma) gemv_mma; +#define MM_SET(F) MM(F, 1) MM(F, 2) MM(F, 4) +MM_SET(T_BF16) MM_SET(T_Q8_0) MM_SET(T_Q4_0) MM_SET(T_Q5_0) MM_SET(T_IQ4_NL) MM_SET(T_IQ4_XS) MM_SET(T_Q2_0) +MM_SET(T_IQ3_S) MM_SET(T_Q3_K) MM_SET(T_Q4_K) MM_SET(T_Q5_K) MM_SET(T_Q6_K) diff --git a/edge0-strata/metal/kernels/moe.metal b/edge0-strata/metal/kernels/moe.metal new file mode 100644 index 00000000..d02961f3 --- /dev/null +++ b/edge0-strata/metal/kernels/moe.metal @@ -0,0 +1,342 @@ +// metal/kernels/moe.metal - routed experts of one layer for a verify window of 1..8 tokens, grouped by expert: +// +// moe_plan one threadgroup: the window's top-k ids -> groups (distinct experts, by id) with their entries +// (token, weight), each (token, k)'s entry index, and the indirect grid sizes of the next two kernels. +// Deterministic (no atomics). +// moe_gu per group, a block of FFN rows: gate and up rows read once, applied to every entry of the group +// (the shared input x is read from cache), h = silu(gate) * up per entry. +// moe_down per group, a block of output rows: ye[entry] = weight * (down . h[entry]). +// moe_combine y[token] = sum over k of ye[entry(token, k)], in k order. +// +// The plan is one int buffer (see P_* below). Expert addressing: MoeArgs (a pack blob or the GGUF tensors): +// [gate rows | up rows | down rows] in GGUF blocks (tools/iq_pack.py's experts.bin layout). + +#define P_NG 0 +#define P_IGU 4 +#define P_IDN 8 +#define P_MAXE 128 +#define P_GEXP 16 +#define P_GSTART (P_GEXP + P_MAXE) +#define P_ETOK (P_GSTART + P_MAXE + 4) +#define P_EW (P_ETOK + P_MAXE) +#define P_TOKENT (P_EW + P_MAXE) +#define P_INTS (P_TOKENT + P_MAXE) + +// Expert e's matrices: gate at W + gate_off + e * gu_stride, up at W + up_off + e * gu_stride, down at +// W + down_off + e * d_stride (a pack's blobs: gate_off 0, strides = the blob; a GGUF: the three tensors). +struct MoeArgs { + ulong gate_off; + ulong up_off; + ulong down_off; + ulong gu_stride; + ulong d_stride; + uint gu_row; // bytes per gate/up row + uint d_row; // bytes per down row + uint n_embd; + uint n_ff; + uint n_tok; + uint topk; + uint n_expert; + uint gx_gu; // threadgroups along the rows for moe_gu (n_ff / rows per threadgroup) + uint gx_dn; // ... for moe_down (n_embd / rows per threadgroup) + uint pad; +}; + +// The SSD-streaming arena (metal/src/engine/stream.cpp): with function constant 20 set, expert e of the layer lies in +// an arena slot, smap[e] = (chunk << 28) | its 16 KiB page in that chunk; the chunks are W (buffer 0) and buffers 7..9 +// (gate_off / up_off / down_off are then offsets inside the slot). The same block is in moe_mm.metal. +#ifndef STRATA_ARENA_METAL +#define STRATA_ARENA_METAL +constant bool kArenaFc [[function_constant(20)]]; +constant bool kArena = is_function_constant_defined(kArenaFc) && kArenaFc; +static inline device const uchar* arena_slot(uint v, device const uchar* W, device const uchar* W1, device const uchar* W2, + device const uchar* W3) { + const uint c = v >> 28; + device const uchar* b = c == 0 ? W : (c == 1 ? W1 : (c == 2 ? W2 : W3)); + return b + (ulong)(v & 0x0FFFFFFFu) * 16384ul; +} +#endif + +static inline void plan_body(device const int* ids, device const float* w, device int* plan, constant MoeArgs& a, + ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp, + threadgroup int* s_ent) { + const int ne = int(a.n_tok * a.topk); + // one thread per expert id (nth == n_expert, at most 1024) + const int e = tid; + int cnt = 0; + if (e < int(a.n_expert)) + for (int i = 0; i < ne; ++i) cnt += ids[i] == e; + const int has = cnt > 0; + const int pg = simd_prefix_exclusive_sum(has), pe = simd_prefix_exclusive_sum(cnt); + if (lane == 31) { s_grp[sg] = pg + has; s_ent[sg] = pe + cnt; } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + int og = 0, oe = 0; + const int nsg = (nth + 31) / 32; + for (int i = 0; i < nsg; ++i) { + const int g = s_grp[i], n = s_ent[i]; + s_grp[i] = og; + s_ent[i] = oe; + og += g; + oe += n; + } + plan[P_NG] = og; + plan[P_IGU] = int(a.gx_gu); plan[P_IGU + 1] = og; plan[P_IGU + 2] = 1; + plan[P_IDN] = int(a.gx_dn); plan[P_IDN + 1] = og; plan[P_IDN + 2] = 1; + plan[P_GSTART + og] = oe; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (has) { + const int gi = s_grp[sg] + pg, off = s_ent[sg] + pe; + plan[P_GEXP + gi] = e; + plan[P_GSTART + gi] = off; + int j = 0; + for (int i = 0; i < ne; ++i) + if (ids[i] == e) { + plan[P_ETOK + off + j] = i / int(a.topk); + plan[P_EW + off + j] = as_type(w[i]); + plan[P_TOKENT + i] = off + j; + ++j; + } + } +} + +kernel void moe_plan(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]], + device int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]], + ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup int s_grp[32], s_ent[32]; + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); +} + +// The router's tail for the window, then the plan, in one threadgroup of n_expert (512) threads: simdgroup t routes +// token t exactly as native_router.cu (pinned llama.cpp topk-moe): softmax over the 512 logits, top 10 by +// probability (equal values: the lower expert index), weights = p / max(sum of the selected, 2^-14), the sum taken +// by winning lane in selection order and reduced with the same xor tree. +constant bool MOE_ROUTE_ONLY [[function_constant(28)]]; +kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], + device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]], + ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup int s_grp[32], s_ent[32]; + if (sg < a.n_tok) { + device const float* l = logits + sg * 512; + float v[16]; + for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; + float mx = -INFINITY; + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + if ((ex & 31) == lane) { + v[ex / 32] = -INFINITY; + ids[sg * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } + if (!is_function_constant_defined(MOE_ROUTE_ONLY) || !MOE_ROUTE_ONLY) { + threadgroup_barrier(mem_flags::mem_device); + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); + } +} + +// ---- gate + up + SwiGLU ------------------------------------------------------------------------------------------- +// One pass over a group's rows for all its NE entries (1..8: an expert appears once per token); NE is uniform per +// threadgroup, so the kernel switches once into a body specialised for it. +template +static inline void gu_body(device const uchar* g0, device const uchar* u0, device const float* x, + device const float* xs16, device float* h, device const int* etok, int c, int row0, int li, + constant MoeArgs& a, threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ag[NR][NE], au[NR][NE]; + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NE; ++t) ag[r][t] = au[r][t] = 0.0f; + const int U = int(a.n_embd) / UNIT; + for (int u = li; u < U; u += LPR) { + float qg[NR][UNIT], qu[NR][UNIT]; + float4 cg[NR], cu[NR]; + for (int r = 0; r < NR; ++r) { + Dec::run(g0 + r * a.gu_row, u, qg[r], cg[r], tab); + Dec::run(u0 + r * a.gu_row, u, qu[r], cu[r], tab); + } + for (int t = 0; t < NE; ++t) { + const int tok = etok[t]; + float xs[UNIT], xsum[2]; + X::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum); + for (int r = 0; r < NR; ++r) { + ag[r][t] += unit_combine(qg[r], cg[r], xs, xsum); + au[r][t] += unit_combine(qu[r], cu[r], xs, xsum); + } + } + } + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NE; ++t) { + const float gv = lane_sum(ag[r][t]), uv = lane_sum(au[r][t]); + if (li == 0) h[(c + t) * int(a.n_ff) + row0 + r] = gv / (1.0f + exp(-gv)) * uv; + } +} + +// bench probe (function constant 22): 1 = no table load (wrong values for the i-quants) +constant int MOE_PROBE [[function_constant(22)]]; +constant int moe_probe = is_function_constant_defined(MOE_PROBE) ? MOE_PROBE : 0; + +template +kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* h [[buffer(2)]], device const int* plan [[buffer(3)]], + constant MoeArgs& a [[buffer(4)]], device const float* xs16 [[buffer(5)]], + device const uint* smap [[buffer(6)]], device const uchar* W1 [[buffer(7)]], + device const uchar* W2 [[buffer(8)]], device const uchar* W3 [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + if (moe_probe != 1) { + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const int g = int(tg.y); + if (g >= plan[P_NG]) return; + const int li = lane % LPR; + const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; + const ulong ex = (ulong)plan[P_GEXP + g]; + const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; + device const int* etok = plan + P_ETOK + c; + device const uchar* E = kArena ? arena_slot(smap[ex], W, W1, W2, W3) : W + ex * a.gu_stride; + device const uchar* g0 = E + a.gate_off + (ulong)row0 * a.gu_row; + device const uchar* u0 = E + a.up_off + (ulong)row0 * a.gu_row; + switch (ne) { + case 1: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 2: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 3: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + case 4: gu_body(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break; + default: + for (int k = 0; k < ne; k += 4) { + const int m = min(4, ne - k); + if (m == 4) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else if (m == 3) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else if (m == 2) gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + else gu_body(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab); + } + } +} + +// ---- down --------------------------------------------------------------------------------------------------------- +template +static inline void dn_body(device const uchar* d0, device const float* h, device const float* hs16, device float* ye, + device const int* plan, int c, int row0, int li, constant MoeArgs& a, + threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT; + float ac[NR][NE]; + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NE; ++t) ac[r][t] = 0.0f; + const int U = int(a.n_ff) / UNIT; + for (int u = li; u < U; u += LPR) { + float q[NR][UNIT]; + float4 sc[NR]; + for (int r = 0; r < NR; ++r) Dec::run(d0 + r * a.d_row, u, q[r], sc[r], tab); + for (int t = 0; t < NE; ++t) { + float xs[UNIT], xsum[2]; + X::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum); + for (int r = 0; r < NR; ++r) ac[r][t] += unit_combine(q[r], sc[r], xs, xsum); + } + } + for (int r = 0; r < NR; ++r) + for (int t = 0; t < NE; ++t) { + const float v = lane_sum(ac[r][t]); + if (li == 0) ye[(c + t) * int(a.n_embd) + row0 + r] = as_type(plan[P_EW + c + t]) * v; + } +} + +template +kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h [[buffer(1)]], + device float* ye [[buffer(2)]], device const int* plan [[buffer(3)]], + constant MoeArgs& a [[buffer(4)]], device const float* hs16 [[buffer(5)]], + device const uint* smap [[buffer(6)]], device const uchar* W1 [[buffer(7)]], + device const uchar* W2 [[buffer(8)]], device const uchar* W3 [[buffer(9)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int RG = 32 / LPR; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int g = int(tg.y); + if (g >= plan[P_NG]) return; + const int li = lane % LPR; + const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR; + const ulong ex = (ulong)plan[P_GEXP + g]; + const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c; + device const uchar* d0 = (kArena ? arena_slot(smap[ex], W, W1, W2, W3) : W + ex * a.d_stride) + a.down_off + (ulong)row0 * a.d_row; + switch (ne) { + case 1: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 2: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 3: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + case 4: dn_body(d0, h, hs16, ye, plan, c, row0, li, a, tab); break; + default: + for (int k = 0; k < ne; k += 4) { + const int m = min(4, ne - k); + if (m == 4) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else if (m == 3) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else if (m == 2) dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + else dn_body(d0, h, hs16, ye, plan, c + k, row0, li, a, tab); + } + } +} + +kernel void moe_combine(device const float* ye [[buffer(0)]], device float* y [[buffer(1)]], + device const int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.n_embd, k = gid % a.n_embd; + if (t >= a.n_tok) return; + float s = 0.0f; + for (uint j = 0; j < a.topk; ++j) s += ye[plan[P_TOKENT + t * a.topk + j] * int(a.n_embd) + k]; + y[gid] = s; +} + +// ---- instantiations: moe_gu____, moe_down____ ------------------------------- +#define GU(F, L, R, S) template [[host_name("moe_gu_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_gu) moe_gu; +#define DN(F, L, R, S) template [[host_name("moe_down_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_down) moe_down; +#define GU_SET(F) GU(F, 32, 1, 4) GU(F, 32, 2, 4) GU(F, 32, 4, 4) GU(F, 32, 2, 2) GU(F, 32, 2, 8) GU(F, 32, 1, 8) \ + GU(F, 16, 1, 4) GU(F, 16, 2, 4) GU(F, 8, 1, 4) GU(F, 8, 2, 4) GU(F, 8, 1, 2) +#define DN_SET(F) DN(F, 8, 1, 4) DN(F, 8, 2, 4) DN(F, 8, 4, 4) DN(F, 8, 1, 8) DN(F, 8, 2, 8) DN(F, 4, 1, 4) DN(F, 4, 2, 4) \ + DN(F, 4, 4, 4) DN(F, 16, 1, 4) DN(F, 16, 2, 4) DN(F, 32, 1, 4) DN(F, 2, 2, 4) +GU_SET(T_Q2_0) +GU_SET(T_IQ3_XXS) +GU_SET(T_IQ2_XXS) +GU_SET(T_IQ2_XS) +GU_SET(T_IQ2_S) +GU_SET(T_IQ3_S) +GU_SET(T_IQ4_XS) // the IQ3_S pack's gate/up in one layer +GU_SET(T_Q4_K) GU_SET(T_Q5_K) GU_SET(T_Q6_K) GU_SET(T_Q8_0) // the UD-Q4/5/6_K_XL packs +DN_SET(T_Q5_1) DN_SET(T_Q8_0) +DN_SET(T_Q2_0) +DN_SET(T_Q2_0R) +DN_SET(T_IQ4_NL) +// Gyro (agentionai's trellis experts): TQK6 / TQ2_T gate/up, TQK7 / TQ2_T down (the CUDA engine's STRATA_GU_FMTS / D_FMTS) +GU_SET(T_TQK6) GU_SET(T_TQ2T) +DN_SET(T_TQK7) DN_SET(T_TQ2T) diff --git a/edge0-strata/metal/kernels/moe_mm.metal b/edge0-strata/metal/kernels/moe_mm.metal new file mode 100644 index 00000000..b4bd0cdd --- /dev/null +++ b/edge0-strata/metal/kernels/moe_mm.metal @@ -0,0 +1,1160 @@ +// metal/kernels/moe_mm.metal - routed experts for a PROMPT chunk (after quant/units/mma.metal): a grouped GEMM on +// the matrix units through Metal 4's MetalPerformancePrimitives matmul2d (on the M5 its neural accelerators). +// +// Plan (deterministic, three passes): the chunk's T x topk (token, expert) entries are sorted by expert, then by +// entry index; each expert's run is padded to a multiple of TM rows; tile i = (expert, first padded row). +// mm_hist per block of 1024 entries: a histogram of experts -> bcnt[block][expert] +// mm_scan one threadgroup (thread = expert): block offsets, the padded expert offsets, the tile list +// mm_place per block: an entry's padded row = expert offset + earlier blocks' count + rank in its block; +// srt_tok[row] = token (-1 on padding), srt_w[row] = weight, pos[entry] = row +// GEMMs: a threadgroup takes a tile (TM rows of one expert) and TN output rows; per K chunk of TK it converts the +// tile's activations to fp16 and dequantizes TN x TK weights (ref_quant's arithmetic, then rounded to fp16) into +// threadgroup memory, and matmul2d multiplies them, fp32 accumulation. +// mm_gu H[row][f] = silu(x . gate_f) * (x . up_f), fp32, for the n_ff rows +// mm_down YE[row][d] = h . down_d, fp32 +// mm_combine y[t][d] = sum_k w(t, k) * YE[pos(t, k)][d], k in order +// +// The idea of feeding matmul2d tiles of dequantized weights from threadgroup memory follows yibie/strata-mlx's +// tile_matmul (MIT); the code is our own. +#include +using namespace mpp::tensor_ops; + +struct MmPlanArgs { + uint n_ent; // T * topk + uint topk; + uint n_expert; + uint tm; // the tile height (rows of one expert) + uint n_blocks; // ceil(n_ent / 1024) +}; + +// counts per block: bcnt[block * n_expert + e] +kernel void mm_hist(device const int* ids [[buffer(0)]], device uint* bcnt [[buffer(1)]], + constant MmPlanArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup atomic_uint h[1024]; + for (uint e = tid; e < a.n_expert; e += 1024) atomic_store_explicit(&h[e], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint i = tg.x * 1024 + tid; + if (i < a.n_ent) atomic_fetch_add_explicit(&h[ids[i]], 1u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint e = tid; e < a.n_expert; e += 1024) + bcnt[tg.x * a.n_expert + e] = atomic_load_explicit(&h[e], memory_order_relaxed); +} + +// one threadgroup of n_expert (<= 1024) threads. boff[block][e] = expert e's rows before this block (within e); +// eoff[e] = expert e's first padded row; tiles[i] = (expert, row0); meta[0] = n_tiles, meta[1] = padded rows; +// indirect grids: meta[2..4] = (gx_gu, n_tiles, 1) for the gate/up GEMM, meta[5..7] = (gx_down, n_tiles, 1) for the +// down GEMM (gx_gu in meta[2], gx_down in meta[5] written by the host once), meta[8..10] = (rows, 1, 1) for mm_gather. +kernel void mm_scan(device const uint* bcnt [[buffer(0)]], device uint* boff [[buffer(1)]], + device uint* eoff [[buffer(2)]], device int2* tiles [[buffer(3)]], device uint* meta [[buffer(4)]], + constant MmPlanArgs& a [[buffer(5)]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup uint s_rows[32], s_tiles[32]; + // simdgroups past the launched ones must count 0 (512 threads fill only 16 of the 32 slots) + if (tid < 32) { s_rows[tid] = 0; s_tiles[tid] = 0; } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint e = tid; + uint tot = 0; + if (e < a.n_expert) + for (uint b = 0; b < a.n_blocks; ++b) { + boff[b * a.n_expert + e] = tot; + tot += bcnt[b * a.n_expert + e]; + } + const uint nt = (tot + a.tm - 1) / a.tm, rows = nt * a.tm; + const uint pr = simd_prefix_exclusive_sum(rows), pt = simd_prefix_exclusive_sum(nt); + if (lane == 31) { s_rows[sg] = pr + rows; s_tiles[sg] = pt + nt; } + threadgroup_barrier(mem_flags::mem_threadgroup); + uint orow = 0, otile = 0, all_rows = 0, all_tiles = 0; + for (int i = 0; i < 32; ++i) { + if (i < int(sg)) { orow += s_rows[i]; otile += s_tiles[i]; } + all_rows += s_rows[i]; + all_tiles += s_tiles[i]; + } + if (e < a.n_expert) { + const uint r0 = orow + pr, t0 = otile + pt; + eoff[e] = r0; + for (uint k = 0; k < nt; ++k) tiles[t0 + k] = int2(int(e), int(r0 + k * a.tm)); + } + if (tid == 0) { + meta[0] = all_tiles; + meta[1] = all_rows; + meta[3] = all_tiles; // indirect grid of the gate/up GEMM: (meta[2] given, tiles, 1) + meta[4] = 1; + meta[6] = all_tiles; // of the down GEMM: (meta[5] given, tiles, 1) + meta[7] = 1; + meta[8] = all_rows; // of the gather: (rows, 1, 1) + meta[9] = 1; + meta[10] = 1; + } +} + +// grid (n_blocks), 1024 threads. srt_* must be cleared to the padding value (-1 / 0) beforehand by mm_clear. +kernel void mm_place(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]], + device const uint* boff [[buffer(2)]], device const uint* eoff [[buffer(3)]], + device int* srt_tok [[buffer(4)]], device float* srt_w [[buffer(5)]], device int* pos [[buffer(6)]], + constant MmPlanArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup int s_id[1024]; + const uint i = tg.x * 1024 + tid; + const int e = i < a.n_ent ? ids[i] : -1; + s_id[tid] = e; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (i >= a.n_ent) return; + uint rank = 0; + for (uint j = 0; j < tid; ++j) rank += s_id[j] == e; + const uint row = eoff[e] + boff[tg.x * a.n_expert + e] + rank; + srt_tok[row] = int(i / a.topk); + srt_w[row] = w[i]; + pos[i] = int(row); +} + +kernel void mm_clear(device int* srt_tok [[buffer(0)]], constant uint& n [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + if (gid < n) srt_tok[gid] = -1; +} + +// ---- the GEMMs ----------------------------------------------------------------------------------------------------- +// weights of one expert: gate at W + gate_off + e * gu_stride (+ up_off for up), down at W + down_off + e * d_stride +struct MmArgs { + ulong gate_off; + ulong up_off; + ulong down_off; + ulong gu_stride; + ulong d_stride; + uint gu_row; // bytes per gate/up row (n_embd values) + uint d_row; // bytes per down row (n_ff values) + uint n_embd; + uint n_ff; +}; + +// TN x TK weights of `nrow`-strided rows into threadgroup half: thread-strided units of the format +template +static inline void stage_w(device const uchar* w0, uint row_bytes, int k0, threadgroup half* dst, + threadgroup const uint* tab, ushort tid, ushort nth) { + constexpr int UNIT = fmt::UNIT, UPC = TK / UNIT, H = UNIT / 2; + static_assert(TK % UNIT == 0, "the K chunk must hold whole units of the format"); + for (int it = tid; it < TN * UPC; it += nth) { + const int r = it / UPC, uu = it - r * UPC; + float wv[UNIT]; + int p0, p1; + DW::run(w0 + ulong(r) * row_bytes, k0 / UNIT + uu, k0, wv, p0, p1, tab); + for (int j = 0; j < H; ++j) { + dst[r * TK + p0 + j] = half(wv[j]); + dst[r * TK + p1 + j] = half(wv[H + j]); + } + } +} + +template +kernel void mm_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]], + device float* H [[buffer(2)]], device const int* srt_tok [[buffer(3)]], + device const int2* tiles [[buffer(4)]], constant MmArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half xa[TM * TK]; + threadgroup half wg[TN * TK]; + threadgroup half wu[TN * TK]; + threadgroup uint tab[fmt::TAB]; + threadgroup int tok[TM]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, f0 = int(tg.x) * TN; + for (int m = tid; m < TM; m += NTH) tok[m] = srt_tok[row0 + m]; + device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(f0) * a.gu_row; + device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(f0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false); + matmul2d> op; + tensor, tensor_inline> tA(xa, dextents(TK, TM)); + tensor, tensor_inline> tG(wg, dextents(TK, TN)); + tensor, tensor_inline> tU(wu, dextents(TK, TN)); + auto accg = op.template get_destination_cooperative_tensor(); + auto accu = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + for (int it = tid; it < TM * TK; it += NTH) { + const int m = it / TK, k = it - m * TK; + xa[it] = tok[m] >= 0 ? half(x[ulong(tok[m]) * a.n_embd + k0 + k]) : half(0.0f); + } + stage_w(g0, a.gu_row, k0, wg, tab, tid, NTH); + stage_w(u0, a.gu_row, k0, wu, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + auto pg = op.template get_destination_cooperative_tensor(); + auto pu = op.template get_destination_cooperative_tensor(); + op.run(tA, tG, pg); + op.run(tA, tU, pu); + for (uint16_t i = 0; i < pg.get_capacity(); ++i) + if (pg.is_valid_element(i)) { accg[i] += pg[i]; accu[i] += pu[i]; } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tH(H, dextents(a.n_ff, 1 << 30)); + accg.store(tH.slice(f0, row0)); +} + +template +kernel void mm_down(device const uchar* W [[buffer(0)]], device const float* H [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], + constant MmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half ha[TM * TK]; + threadgroup half wd[TN * TK]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, d0 = int(tg.x) * TN; + device const uchar* w0 = W + a.down_off + ulong(e) * a.d_stride + ulong(d0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false); + matmul2d> op; + tensor, tensor_inline> tA(ha, dextents(TK, TM)); + tensor, tensor_inline> tW(wd, dextents(TK, TN)); + auto acc = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + for (int it = tid; it < TM * TK; it += NTH) { + const int m = it / TK, k = it - m * TK; + ha[it] = half(H[ulong(row0 + m) * a.n_ff + k0 + k]); + } + stage_w(w0, a.d_row, k0, wd, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + auto p = op.template get_destination_cooperative_tensor(); + op.run(tA, tW, p); + for (uint16_t i = 0; i < p.get_capacity(); ++i) + if (p.is_valid_element(i)) acc[i] += p[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(d0, row0)); +} + +// grid (T * n_embd / 256), 256 threads +kernel void mm_combine(device const float* YE [[buffer(0)]], device const int* pos [[buffer(1)]], + device const float* srt_w [[buffer(2)]], device float* y [[buffer(3)]], + constant uint4& a [[buffer(4)]], // (T, n_embd, topk, -) + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.y, d = gid % a.y; + if (t >= a.x) return; + float s = 0.0f; + for (uint k = 0; k < a.z; ++k) { + const int r = pos[t * a.z + k]; + s = fma(srt_w[r], YE[ulong(r) * a.y + d], s); + } + y[gid] = s; +} + +#define MMGU_(F, M, N, K, S) template [[host_name("mm_gu_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu) mm_gu; +#define MMDN_(F, M, N, K, S) template [[host_name("mm_down_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down) mm_down; +#define MMGU(F, M, N, K, S) MMGU_(F, M, N, K, S) +#define MMDN(F, M, N, K, S) MMDN_(F, M, N, K, S) +#define MM_GU_SET(F) MMGU(F, 32, 64, 64, 4) MMGU(F, 64, 64, 64, 4) MMGU(F, 16, 64, 64, 2) MMGU(F, 32, 32, 64, 2) +#define MM_DN_SET(F) MMDN(F, 32, 64, 64, 4) MMDN(F, 64, 64, 64, 4) MMDN(F, 16, 64, 64, 2) MMDN(F, 32, 32, 64, 2) +MM_GU_SET(T_Q2_0) MM_GU_SET(T_IQ3_XXS) MM_GU_SET(T_IQ2_XXS) MM_GU_SET(T_IQ2_XS) MM_GU_SET(T_IQ2_S) MM_GU_SET(T_IQ3_S) +MM_DN_SET(T_Q2_0) MM_DN_SET(T_IQ4_NL) + +// ---- v2: activations gathered once per layer as fp16 rows, read by matmul2d from device memory --------------------- +// mm_gather: XS[row][k] = half(x[srt_tok[row]][k]) (0 on padding rows). grid (rows), 256 threads. +kernel void mm_gather(device const float* x [[buffer(0)]], device const int* srt_tok [[buffer(1)]], + device half* XS [[buffer(2)]], constant uint& n [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const int t = srt_tok[tg.x]; + device half* o = XS + ulong(tg.x) * n; + for (uint k = tid; k < n; k += 256) o[k] = t >= 0 ? half(x[ulong(t) * n + k]) : half(0.0f); +} + +// The SSD-streaming arena (as in moe.metal): with function constant 20, expert e's slot is smap[e] in the chunks W, +// W1..W3 (buffers 5..8 of mm_gu2 / mm_down2 / mm_gu3 / mm_down3); gate_off / up_off / down_off are offsets in the slot. +#ifndef STRATA_ARENA_METAL +#define STRATA_ARENA_METAL +constant bool kArenaFc [[function_constant(20)]]; +constant bool kArena = is_function_constant_defined(kArenaFc) && kArenaFc; +static inline device const uchar* arena_slot(uint v, device const uchar* W, device const uchar* W1, device const uchar* W2, + device const uchar* W3) { + const uint c = v >> 28; + device const uchar* b = c == 0 ? W : (c == 1 ? W1 : (c == 2 ? W2 : W3)); + return b + (ulong)(v & 0x0FFFFFFFu) * 16384ul; +} +#endif + +// a probe for the bench (function constant 5): 1 = the weights not decoded (a cheap fill), 2 = no matrix multiply +constant int MM_PROBE [[function_constant(5)]]; +constant int mm_probe = is_function_constant_defined(MM_PROBE) ? MM_PROBE : 0; + +template +kernel void mm_gu2(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half wg[TN * TK]; + threadgroup half wu[TN * TK]; + threadgroup float ot[TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, f0 = int(tg.x) * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* g0 = E + a.gate_off + ulong(f0) * a.gu_row; + device const uchar* u0 = E + a.up_off + ulong(f0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + tensor, tensor_inline> tG(wg, dextents(TK, TN)); + tensor, tensor_inline> tU(wu, dextents(TK, TN)); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor(); + auto accu = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + if (mm_probe == 1) { + for (int it = tid; it < TN * TK; it += NTH) { wg[it] = half(k0 & 7); wu[it] = half(it & 7); } + } else { + stage_w(g0, a.gu_row, k0, wg, tab, tid, NTH); + stage_w(u0, a.gu_row, k0, wu, tab, tid, NTH); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + tA = tX.template slice(k0, row0); + if (mm_probe != 2) { + op.run(tA, tG, accg); + op.run(tA, tU, accu); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot, dextents(TN, TM)); + accg.store(tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int it = tid; it < TM * TN; it += NTH) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + f0 + f] = half(ot[it]); + } +} + +template +kernel void mm_down2(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], + constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int NTH = SG * 32; + threadgroup half wd[TN * TK]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, NTH); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, d0 = int(tg.x) * TN; + device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(d0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d> op; + tensor, tensor_inline> tH((device half*)HS, dextents(a.n_ff, 1 << 30)); + tensor, tensor_inline> tW(wd, dextents(TK, TN)); + auto tA = tH.template slice(0, row0); + auto acc = op.template get_destination_cooperative_tensor(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + stage_w(w0, a.d_row, k0, wd, tab, tid, NTH); + threadgroup_barrier(mem_flags::mem_threadgroup); + tA = tH.template slice(k0, row0); + op.run(tA, tW, acc); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(d0, row0)); +} + +#define MMGU2_(F, M, N, K, S) template [[host_name("mm_gu2_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu2) mm_gu2; +#define MMDN2_(F, M, N, K, S) template [[host_name("mm_down2_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down2) mm_down2; +#define MMGU2(F, M, N, K, S) MMGU2_(F, M, N, K, S) +#define MMDN2(F, M, N, K, S) MMDN2_(F, M, N, K, S) +#define MM_GU2_SET(F) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 32, 32, 32, 2) MMGU2(F, 32, 64, 16, 2) MMGU2(F, 64, 32, 16, 2) MMGU2(F, 16, 32, 32, 2) MMGU2(F, 16, 64, 16, 2) MMGU2(F, 16, 128, 32, 4) MMGU2(F, 32, 64, 16, 4) MMGU2(F, 16, 32, 32, 1) MMGU2(F, 16, 64, 32, 4) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 64, 4) MMGU2(F, 64, 64, 64, 4) MMGU2(F, 16, 32, 64, 2) MMGU2(F, 32, 64, 32, 4) MMGU2(F, 64, 64, 32, 4) MMGU2(F, 16, 64, 32, 2) +#define MM_DN2_SET(F) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 32, 32, 32, 2) MMDN2(F, 32, 64, 16, 2) MMDN2(F, 64, 32, 16, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 16, 64, 16, 2) MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 16, 4) MMDN2(F, 16, 32, 32, 1) MMDN2(F, 16, 64, 32, 4) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 32, 64, 64, 4) MMDN2(F, 64, 64, 64, 4) MMDN2(F, 16, 32, 64, 2) MMDN2(F, 32, 64, 32, 4) MMDN2(F, 64, 64, 32, 4) MMDN2(F, 16, 64, 32, 2) +#define MM_GU2_U32(F) MMGU2(F, 16, 64, 32, 2) MMGU2(F, 16, 64, 64, 2) MMGU2(F, 32, 64, 32, 2) MMGU2(F, 16, 32, 32, 2) \ + MMGU2(F, 64, 64, 32, 4) +#define MM_DN2_U32(F) MMDN2(F, 16, 64, 32, 2) MMDN2(F, 16, 64, 64, 2) MMDN2(F, 16, 32, 32, 2) MMDN2(F, 32, 32, 32, 2) \ + MMDN2(F, 16, 128, 32, 4) MMDN2(F, 32, 64, 32, 2) MMDN2(F, 64, 64, 32, 4) +// K-quants: the staged K chunk must hold whole units as DW places them (128 wide): TK 128 +#define MM_GU2_KQ(F) MMGU2(F, 16, 32, 128, 2) MMGU2(F, 16, 16, 128, 2) MMGU2(F, 32, 32, 128, 2) MMGU2(F, 16, 32, 128, 4) +MM_GU2_KQ(T_Q4_K) MM_GU2_KQ(T_Q5_K) MM_GU2_KQ(T_Q6_K) MM_GU2_U32(T_Q8_0) MM_DN2_U32(T_Q5_1) MM_DN2_U32(T_Q8_0) +// IQ4_XS (UNIT 32): the K chunk must be >= 32 +MMGU2(T_IQ4_XS, 16, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 64, 64, 2) MMGU2(T_IQ4_XS, 32, 64, 32, 2) MMGU2(T_IQ4_XS, 16, 32, 32, 2) +MMGU2(T_IQ4_XS, 64, 64, 32, 4) +MM_GU2_SET(T_Q2_0) MM_GU2_SET(T_IQ3_XXS) MM_GU2_SET(T_IQ2_XXS) MM_GU2_SET(T_IQ2_XS) MM_GU2_SET(T_IQ2_S) MM_GU2_SET(T_IQ3_S) +MMGU2(T_Q3_K, 16, 64, 64, 2) +MMGU2(T_Q3_K, 32, 64, 32, 2) +MMDN2(T_IQ4_NL, 32, 64, 32, 2) +MM_GU2_SET(T_TQK6) MM_GU2_SET(T_TQ2T) // Gyro: the trellis gate/up (units of 16 weights: TK 16 or a multiple) +MM_DN2_SET(T_Q2_0) +MM_DN2_SET(T_TQK7) MM_DN2_SET(T_TQ2T) +MMDN2(T_IQ4_NL, 16, 64, 32, 2) MMDN2(T_IQ4_NL, 16, 64, 64, 2) MMDN2(T_IQ4_NL, 16, 32, 32, 2) MMDN2(T_IQ4_NL, 32, 32, 32, 2) MMDN2(T_IQ4_NL, 16, 128, 32, 4) + +// ---- v3: the weights decoded straight into matmul2d's right-input cooperative tensor (single-simdgroup scope), +// no threadgroup staging and no barriers in the K loop. Each simdgroup owns TN output rows; a threadgroup SG * TN. +// The right input's element layout (tests/coop_layout): runs of 4 consecutive k at one n, so a lane decodes 4 values +// at a time (W4). Activations: the gathered fp16 rows from device memory. +#ifndef STRATA_COOP_K4 +#define STRATA_COOP_K4 +// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's +// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own. +constant bool COOP_K4 [[function_constant(7)]]; +constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true; +#endif +template struct W4; +// the four weights k .. k + 3 (k % 4 == 0) of a row, in quant.metal's deq arithmetic (code * scale in fp32), as fp16 +// the threadgroup forms of quant.metal's field readers (blocks staged in threadgroup memory) +static inline uint rd32(threadgroup const uchar* p) { + threadgroup const ushort* s = (threadgroup const ushort*)p; + return uint(s[0]) | (uint(s[1]) << 16); +} +static inline float rdh(threadgroup const uchar* p) { return float(*(threadgroup const half*)p); } + +static inline half4 w4_signed(uint g, uint s, float sc) { + const uchar4 v = as_type(g); + const float4 q = float4((s & 1) ? -float(v[0]) : float(v[0]), (s & 2) ? -float(v[1]) : float(v[1]), + (s & 4) ? -float(v[2]) : float(v[2]), (s & 8) ? -float(v[3]) : float(v[3])); + return half4(q * sc); +} +template <> struct W4 { // 64 values a block: fp16 d, 16 bytes of 2-bit codes (value v at byte v / 4) + // fp32 math then fp16 (an fp16-math decode was measured slower: 4.2 vs 3.5 ms a layer) + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 6) * 18; + const float d = float(*(device const half*)b); + const uint q = b[2 + ((k & 63) >> 2)]; + return half4(float4(float(int(q & 3) - 1), float(int((q >> 2) & 3) - 1), float(int((q >> 4) & 3) - 1), + float(int(q >> 6) - 1)) * d); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 98; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint aux = rd32(b + 2 + 64 + 4 * ib); + threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 256); + const uint s = ks[(aux >> (7 * l)) & 127]; + const float sc = rdh(b) * (0.5f + float(aux >> 28)) * 0.5f; + return w4_signed(tab[b[2 + 8 * ib + 2 * l + hh]], hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 66; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4); + threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 512); + const uint idx = (a0 >> (8 * l)) & 255, s = ks[(a1 >> (7 * l)) & 127]; + const float sc = rdh(b) * (0.5f + float(a1 >> 28)) * 0.25f; + return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 74; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint v = ((device const ushort*)(b + 2))[4 * ib + l]; + threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 1024); + const uint idx = v & 511, s = ks[v >> 9]; + const uint scb = b[66 + ib]; + const float sc = rdh(b) * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; + return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 82; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint qh = b[66 + ib]; + const uint idx = uint(b[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300); + const uint s = b[2 + 32 + 4 * ib + l]; + const uint scb = b[74 + ib]; + const float sc = rdh(b) * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; + return w4_signed(tab[2 * idx + hh], hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 110; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint qh = b[66 + ib]; + const uint i = hh ? (uint(b[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256)) + : (uint(b[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256)); + const uint s = b[74 + 4 * ib + l]; + const uint scb = b[106 + (ib >> 1)]; + const float sc = rdh(b) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15)); + return w4_signed(tab[i], hh ? s >> 4 : s, sc); + } +}; +template <> struct W4 { // 256 a block: d, scales_h, scales_l[4], qs[128]; 32-groups as IQ4_NL's + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 8) * 136; + const int j = k & 255, ib = j >> 5, o = j & 15, hi = (j >> 4) & 1; + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ib >> 1)] >> (4 * (ib & 1))) & 15) | int(((sh >> (2 * ib)) & 3) << 4); + const float s = rdh(b) * float(ls - 32); + device const uchar* qs = b + 8 + 16 * ib; + float4 q; + for (int c = 0; c < 4; ++c) { + const half2 w = as_type(tab[qs[o + c]]); + q[c] = float(hi ? w.y : w.x); + } + return half4(q * s); + } +}; +template <> struct W4 { // 32 a block: d, 16 bytes; low nibbles are elements 0..15, high 16..31 + static inline half4 run(device const uchar* row, int k, threadgroup const uint* tab) { + device const uchar* b = row + (k >> 5) * 18; + const int j = k & 31, hi = j >> 4, o = j & 15; + const float d = rdh(b); + float4 q; + for (int c = 0; c < 4; ++c) { + const half2 w = as_type(tab[b[2 + o + c]]); + q[c] = float(hi ? w.y : w.x); + } + return half4(q * d); + } +}; + +// W8: the runs k .. k + 3 and k + 16 .. k + 19 of a row (k % 32 < 16): one 32-group, so the per-group fields (scale, +// sign word) are read once. The generic one is two W4. +// ---- the K-quant / legacy-quant W4 decoders (the UD-Q4/5/6_K_XL packs' experts): values k .. k + 3 (k % 4 == 0) +// of a row in ref_quant's arithmetic (float), then fp16 +template <> struct W4 { // 32 a block: d, int8 qs[32] + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 5) * 34; + const float d = rdh(b); + device const char* q = (device const char*)(b + 2 + (k & 31)); + return half4(float4(float(q[0]), float(q[1]), float(q[2]), float(q[3])) * d); + } +}; +template <> struct W4 { // 32 a block: d, m, qh (32 bits), qs[16] (low nibbles 0..15, high 16..31) + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 5) * 24; + const float d = rdh(b), m = rdh(b + 2); + const uint hb = rd32(b + 4); + const int j = k & 31, o = j & 15, hi = j >> 4; + float4 w; + for (int c = 0; c < 4; ++c) { + const uint qv = b[8 + o + c]; + const int x = int(hi ? (qv >> 4) : (qv & 15)) | int(((hb >> (j + c)) & 1) << 4); + w[c] = float(x) * d + m; + } + return half4(w); + } +}; +template <> struct W4 { // 256 a block: d, dmin, scales[12], qs[128] (64-chunks: low nibbles then high) + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 144; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, b + 4, s, mm); + const float d1 = rdh(b) * s, m1 = rdh(b + 2) * mm; + device const uchar* q = b + 16 + 32 * c + l; + float4 w; + for (int i = 0; i < 4; ++i) w[i] = d1 * float(hi ? (q[i] >> 4) : (q[i] & 15)) - m1; + return half4(w); + } +}; +template <> struct W4 { // 256 a block: d, dmin, scales[12], qh[32], qs[128] + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 176; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, b + 4, s, mm); + const float d1 = rdh(b) * s, m1 = rdh(b + 2) * mm; + device const uchar* qh = b + 16 + l; + device const uchar* ql = b + 48 + 32 * c + l; + const uint bit = 1u << (2 * c + hi); + float4 w; + for (int i = 0; i < 4; ++i) + w[i] = d1 * float(int(hi ? (ql[i] >> 4) : (ql[i] & 15)) + ((qh[i] & bit) ? 16 : 0)) - m1; + return half4(w); + } +}; +template <> struct W4 { // 256 a block: ql[128], qh[64], int8 scales[16], d; halves of 128 + static inline half4 run(device const uchar* row, int k, threadgroup const uint*) { + device const uchar* b = row + (k >> 8) * 210; + const int j = k & 255, n = j >> 7, jj = j & 127, qd = jj >> 5, l = jj & 31; + device const uchar* ql = b + 64 * n + l + ((qd & 1) ? 32 : 0); + device const uchar* qh = b + 128 + 32 * n + l; + const float d = rdh(b + 208); + const float ds = d * float(((device const char*)(b + 192 + 8 * n))[l / 16 + 2 * qd]); + float4 w; + for (int i = 0; i < 4; ++i) { + const int q = int(((qd >> 1) ? (ql[i] >> 4) : (ql[i] & 15)) | (((qh[i] >> (2 * qd)) & 3) << 4)) - 32; + w[i] = ds * float(q); + } + return half4(w); + } +}; + +template struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + a = W4::run(row, k, tab); + b = W4::run(row, k + 16, tab); + } +}; +// (a W8 sharing the block scale was measured slower than the two W4 of the generic one: 5.5 vs 3.5-4.2 ms) +template <> struct W8 { + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 98; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint aux = rd32(bl + 2 + 64 + 4 * ib); + threadgroup const uchar* ks = (threadgroup const uchar*)(tab + 256); + const float sc = rdh(bl) * (0.5f + float(aux >> 28)) * 0.5f; + const P qs = bl + 2 + 8 * ib + hh; + if (mm_probe == 3) { // bench probe: no threadgroup-table lookups (wrong values) + a = w4_signed(qs[2 * l] * 0x01010101u, aux, sc); + b = w4_signed(qs[2 * l + 4] * 0x01010101u, aux >> 7, sc); + return; + } + uint s0, s1; + if (mm_probe == 4) { // the signs computed (ksigns_iq2xs[i] = i | parity(i) << 7) instead of looked up + const uint i0 = (aux >> (7 * l)) & 127, i1 = (aux >> (7 * (l + 2))) & 127; + s0 = i0 | ((popcount(i0) & 1u) << 7); + s1 = i1 | ((popcount(i1) & 1u) << 7); + } else { + s0 = ksg(tab, 256, (aux >> (7 * l)) & 127); + s1 = ksg(tab, 256, (aux >> (7 * (l + 2))) & 127); + } + a = w4_signed(g_iq3xxs(tab, qs[2 * l]), hh ? s0 >> 4 : s0, sc); + b = w4_signed(g_iq3xxs(tab, qs[2 * l + 4]), hh ? s1 >> 4 : s1, sc); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 66; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint a0 = rd32(bl + 2 + 8 * ib), a1 = rd32(bl + 2 + 8 * ib + 4); + const float sc = rdh(bl) * (0.5f + float(a1 >> 28)) * 0.25f; + const uint i0 = (a0 >> (8 * l)) & 255, i1 = (a0 >> (8 * (l + 2))) & 255; + // through the accessors, so QUANT_CTAB (the constant grids, no threadgroup copy) gives the same values: the + // direct tab[] reads below skipped the ctab branch and read an unfilled table + const uint s0 = ksg(tab, 512, (a1 >> (7 * l)) & 127), s1 = ksg(tab, 512, (a1 >> (7 * (l + 2))) & 127); + a = w4_signed(g_iq2xxs(tab, i0, hh), hh ? s0 >> 4 : s0, sc); + b = w4_signed(g_iq2xxs(tab, i1, hh), hh ? s1 >> 4 : s1, sc); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 74; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + device const ushort* q = (device const ushort*)(bl + 2) + 4 * ib; + const uint v0 = q[l], v1 = q[l + 2]; + const uint scb = bl[66 + ib]; + const float d = rdh(bl); + const float sc0 = d * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; + const float sc1 = d * (0.5f + float((scb >> (4 * ((l + 2) >> 1))) & 15)) * 0.25f; + // through the accessors: the direct tab[]/ks[] reads ignored QUANT_CTAB, whose tables are then never filled + const uint s0 = ksg(tab, 1024, v0 >> 9), s1 = ksg(tab, 1024, v1 >> 9); + a = w4_signed(g_iq2xs(tab, v0 & 511, hh), hh ? s0 >> 4 : s0, sc0); + b = w4_signed(g_iq2xs(tab, v1 & 511, hh), hh ? s1 >> 4 : s1, sc1); + } +}; +template <> struct W8 { + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 82; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint qh = bl[66 + ib], scb = bl[74 + ib]; + const float d = rdh(bl); + const uint i0 = uint(bl[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300); + const uint i1 = uint(bl[2 + 4 * ib + l + 2]) | ((qh << (8 - 2 * (l + 2))) & 0x300); + const uint s0 = bl[2 + 32 + 4 * ib + l], s1 = bl[2 + 32 + 4 * ib + l + 2]; + const float sc0 = d * (0.5f + float((scb >> (4 * (l >> 1))) & 15)) * 0.25f; + const float sc1 = d * (0.5f + float((scb >> (4 * ((l + 2) >> 1))) & 15)) * 0.25f; + a = w4_signed(g_iq2s(tab, i0, hh), hh ? s0 >> 4 : s0, sc0); + b = w4_signed(g_iq2s(tab, i1, hh), hh ? s1 >> 4 : s1, sc1); + } +}; + +// IQ1_M (the 2.00 bpw pack's 1.75-bit type), ggml's dequantize_row_iq1_m: the block's F16 scale sits in the top +// nibbles of its four u16 scales; a 32-value sub-block has two 3-bit scales (groups 0-1 and 2-3, eight values each), +// and each eight-value group carries a +-0.125 delta picked by a bit of its qh byte. The tile's two halves are the +// four values at k and the four at k + 16, each decoded on its own (k + 16 crosses into the next sub-block at l >= 2). +template +static inline half4 iq1m_w4(P row, int k) { + const P bl = row + (k >> 8) * 56; + const int ib = (k & 255) >> 5, g = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint s0 = uint(bl[48]) | (uint(bl[49]) << 8), s1 = uint(bl[50]) | (uint(bl[51]) << 8); + const uint s2 = uint(bl[52]) | (uint(bl[53]) << 8), s3 = uint(bl[54]) | (uint(bl[55]) << 8); + const float d = float(as_type(ushort((s0 >> 12) | ((s1 >> 8) & 0x00f0u) | ((s2 >> 4) & 0x0f00u) | (s3 & 0xf000u)))); + const uint sc = uint(((const device ushort*) (bl + 48))[ib >> 1]); + const float dl = d * float(2u * ((sc >> (6u * uint(ib & 1) + (g < 2 ? 0u : 3u))) & 7u) + 1u); + const uint qh = (g < 2) ? uint(bl[32 + 2 * ib]) : uint(bl[33 + 2 * ib]); + const uint idx = uint(bl[4 * ib + g]) | (((g & 1) != 0) ? ((qh << 4u) & 0x700u) : ((qh << 8u) & 0x700u)); + const bool neg = ((g & 1) != 0) ? ((qh & 0x80u) != 0u) : ((qh & 0x08u) != 0u); + const ulong gr = iq1s_grid[idx]; + const float delta = neg ? -0.125f : 0.125f; + float v[4]; + for (int j = 0; j < 4; ++j) { + const uint by = uint((gr >> (8u * uint(4 * hh + j))) & 0xffu); + const int c = (by & 0x80u) != 0u ? int(by) - 256 : int(by); + v[j] = dl * (float(c) + delta); + } + return half4(v[0], v[1], v[2], v[3]); +} +template <> struct W8 { + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + (void) tab; + a = iq1m_w4(row, k); + b = iq1m_w4(row, k + 16); + } +}; + +// IQ1_S (Unsloth's UD-IQ1_S / UD-IQ1_M experts), ggml's dequantize_row_iq1_s: f16 d, qs[32], qh[8] u16; a 32-value sub-block has +// one scale dl = d (2 s + 1) (qh bits 12-14), a +-0.125 delta (qh bit 15) and four 8-value groups whose grid index is +// qs[4 ib + g] | 3 bits of qh << 8. Same float arithmetic as ggml (dl (grid + delta)), so the half4 is what the CPU reference rounds to. +template +static inline half4 iq1s_w4(P row, int k) { + const P bl = row + (k >> 8) * 50; + const int ib = (k & 255) >> 5, g = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint qh = uint(bl[34 + 2 * ib]) | (uint(bl[35 + 2 * ib]) << 8); + const float dl = float(as_type(ushort(uint(bl[0]) | (uint(bl[1]) << 8)))) * float(2u * ((qh >> 12) & 7u) + 1u); + const float delta = (qh & 0x8000u) != 0u ? -0.125f : 0.125f; + const uint idx = uint(bl[2 + 4 * ib + g]) | (((qh >> (3u * uint(g))) & 7u) << 8); + const ulong gr = iq1s_grid[idx]; + float v[4]; + for (int j = 0; j < 4; ++j) { + const uint by = uint((gr >> (8u * uint(4 * hh + j))) & 0xffu); + const int c = (by & 0x80u) != 0u ? int(by) - 256 : int(by); + v[j] = dl * (float(c) + delta); + } + return half4(v[0], v[1], v[2], v[3]); +} +template <> struct W8 { + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + (void) tab; + a = iq1s_w4(row, k); + b = iq1s_w4(row, k + 16); + } +}; + +template <> struct W8 { + template + static inline void run(P row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + const P bl = row + (k >> 8) * 110; + const int ib = (k & 255) >> 5, l = (k & 31) >> 3, hh = (k & 7) >> 2; + const uint qh = bl[66 + ib], scb = bl[106 + (ib >> 1)]; + const float sc = rdh(bl) * float(1 + 2 * ((scb >> (4 * (ib & 1))) & 15)); + const int l1 = l + 2; + const uint i0 = hh ? (uint(bl[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256)) + : (uint(bl[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256)); + const uint i1 = hh ? (uint(bl[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256)) + : (uint(bl[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256)); + const uint s0 = bl[74 + 4 * ib + l], s1 = bl[74 + 4 * ib + l1]; + if (mm_probe == 7) { + // bench probe: a repacked layout's cost - one aligned 16-byte load a lane a row a 32-group holding every + // field the lane needs (wrong values: the bytes come from an arbitrary aligned spot of the row) + const uint4 r = *(device const uint4*)((ulong(row) + ulong(k >> 5) * 16) & ~ulong(15)); + const uint qhv = r.z & 0xFF, scv = (r.z >> 8) & 0xFF; + const float scr = float(as_type(ushort(r.w & 0xFFFF))) * float(1 + 2 * ((scv >> (4 * (ib & 1))) & 15)); + const uint j0 = ((r.x >> (8 * (l & 1))) & 0xFF) | ((qhv << (8 - 2 * l)) & 256); + const uint j1 = ((r.x >> (16 + 8 * (l & 1))) & 0xFF) | ((qhv << (7 - 2 * l)) & 256); + a = w4_signed(tab[j0], (r.y >> (8 * hh)) & 0xFF, scr); + b = w4_signed(tab[j1], (r.y >> (16 + 8 * hh)) & 0xFF, scr); + return; + } + if (mm_probe == 5) { // bench probe: the loads and the scale, no grid lookup / sign math (wrong values) + a = half4(half(float(i0 ^ s0) * sc)); + b = half4(half(float(i1 ^ s1) * sc)); + return; + } + a = w4_signed(g_iq3s(tab, i0), hh ? s0 >> 4 : s0, sc); + b = w4_signed(g_iq3s(tab, i1), hh ? s1 >> 4 : s1, sc); + } +}; + +// W8 for the K-quants and Q5_1: the lane's runs k and k + 16 (k % 32 < 16) share the block's scales / mins / high +// bits; the 4-byte code runs read as one uint where the block layout keeps them 4-aligned (Q4_K 144-, Q5_K 176-, +// Q5_1 24-byte blocks: rows of whole blocks start 4-aligned). The same per-value arithmetic as W4. +static inline float4 nib4(uint v, int hi) { const uint s = hi ? (v >> 4) : v; return float4(as_type(s & 0x0F0F0F0Fu)); } +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 144; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, bl + 4, s, mm); + const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm; + device const uint* q = (device const uint*)(bl + 16 + 32 * c + l); + const float4 qa = nib4(q[0], hi), qb = nib4(q[4], hi); + a = half4(d1 * qa - m1); + b = half4(d1 * qb - m1); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 176; + const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31; + float s, mm; + scale_min_k4(2 * c + hi, bl + 4, s, mm); + const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm; + device const uint* qh = (device const uint*)(bl + 16 + l); + device const uint* ql = (device const uint*)(bl + 48 + 32 * c + l); + const uint sh = uint(2 * c + hi); + const float4 ha = float4(as_type((qh[0] >> sh) & 0x01010101u)) * 16.0f; + const float4 hb = float4(as_type((qh[4] >> sh) & 0x01010101u)) * 16.0f; + a = half4(d1 * (nib4(ql[0], hi) + ha) - m1); + b = half4(d1 * (nib4(ql[4], hi) + hb) - m1); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 5) * 24; + const float d = rdh(bl), m = rdh(bl + 2); + const uint hb = rd32(bl + 4); + const int j = k & 15; // runs at j and j + 16: the low nibbles of qs[j..j+3], then the high ones + const uint v = *(device const uint*)(bl + 8 + j); + float4 wa, wb; + for (int c = 0; c < 4; ++c) { + const uint qv = (v >> (8 * c)) & 0xFFu; + wa[c] = float(int(qv & 15) | int(((hb >> (j + c)) & 1) << 4)) * d + m; + wb[c] = float(int(qv >> 4) | int(((hb >> (j + 16 + c)) & 1) << 4)) * d + m; + } + a = half4(wa); + b = half4(wb); + } +}; +template <> struct W8 { + static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) { + device const uchar* bl = row + (k >> 8) * 210; + const int j = k & 255, n = j >> 7, jj = j & 127, qd = jj >> 5, l = jj & 31; // l < 16; the pair at l + 16 + device const uchar* ql = bl + 64 * n + l + ((qd & 1) ? 32 : 0); + device const uchar* qh = bl + 128 + 32 * n + l; + const float d = rdh(bl + 208); + device const char* sc = (device const char*)(bl + 192 + 8 * n); + const float dsa = d * float(sc[2 * qd]), dsb = d * float(sc[1 + 2 * qd]); + // the 4-byte runs as two ushorts each (210-byte blocks: 2-aligned) + device const ushort* l2 = (device const ushort*)ql; + device const ushort* h2 = (device const ushort*)qh; + const uint la = uint(l2[0]) | (uint(l2[1]) << 16), lb = uint(l2[8]) | (uint(l2[9]) << 16); + const uint ha = uint(h2[0]) | (uint(h2[1]) << 16), hb = uint(h2[8]) | (uint(h2[9]) << 16); + const uint sl = (qd >> 1) ? 4u : 0u, shh = uint(2 * qd); + const uint4 qa = (uint4(as_type((la >> sl) & 0x0F0F0F0Fu)) | (uint4(as_type((ha >> shh) & 0x03030303u)) << 4)); + const uint4 qb = (uint4(as_type((lb >> sl) & 0x0F0F0F0Fu)) | (uint4(as_type((hb >> shh) & 0x03030303u)) << 4)); + a = half4(dsa * (float4(qa) - 32.0f)); + b = half4(dsb * (float4(qb) - 32.0f)); + } +}; + +template +kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* E = kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.gu_stride; + device const uchar* g0 = E + a.gate_off + ulong(n0) * a.gu_row; + device const uchar* u0 = E + a.up_off + ulong(n0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rg = op.template get_right_input_cooperative_tensor(); + auto ru = op.template get_right_input_cooperative_tensor(); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + const uint16_t cap = rg.get_capacity(), half_cap = cap / 2; + if (!coop_k4) { + for (uint16_t i = 0; i < cap; ++i) { + const auto ix = rg.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + rg[i] = W4::run(g0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + ru[i] = W4::run(u0 + ulong(n) * a.gu_row, k & ~3, tab)[k & 3]; + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rg.get_multidimensional_index(i); + const auto iy = rg.get_multidimensional_index(i + half_cap); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* rgw = g0 + ulong(n) * a.gu_row; + device const uchar* ruw = u0 + ulong(n) * a.gu_row; + half4 wg, wu, wg2, wu2; + if (mm_probe == 1) { // the bench probe: no decode (constant weights) + wg = half4(half(n & 7)); wu = wg; wg2 = half4(half(k & 7)); wu2 = wg2; + } else if (TK == 32) { + // the lane's two runs of this row sit in one 32-group: element i + cap / 2 is (k + 16, n) - part of + // the COOP_K4 layout (coop_probe checks it); a runtime check of it here cost 30% + W8::run(rgw, k, tab, wg, wg2); + W8::run(ruw, k, tab, wu, wu2); + } else { + const int k2 = k0 + iy[0]; + device const uchar* rgw2 = g0 + ulong(iy[1]) * a.gu_row; + device const uchar* ruw2 = u0 + ulong(iy[1]) * a.gu_row; + wg = W4::run(rgw, k, tab); wu = W4::run(ruw, k, tab); + wg2 = W4::run(rgw2, k2, tab); wu2 = W4::run(ruw2, k2, tab); + } + rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; + ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; + ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + } + tA = tX.template slice(k0, row0); + op.run(tA, rg, accg); + op.run(tA, ru, accu); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + accg.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } +} + +template +kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS [[buffer(1)]], + device float* YE [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], device const uint* smap [[buffer(5)]], + device const uchar* W1 [[buffer(6)]], device const uchar* W2 [[buffer(7)]], device const uchar* W3 [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* w0 = (kArena ? arena_slot(smap[e], W, W1, W2, W3) : W + ulong(e) * a.d_stride) + a.down_off + ulong(n0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tH((device half*)HS, dextents(a.n_ff, 1 << 30)); + auto rw = op.template get_right_input_cooperative_tensor(); + auto tA = tH.template slice(0, row0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + if (!coop_k4) { + for (uint16_t i = 0; i < rw.get_capacity(); ++i) { + const auto ix = rw.get_multidimensional_index(i); + const int k = k0 + ix[0]; + rw[i] = W4::run(w0 + ulong(ix[1]) * a.d_row, k & ~3, tab)[k & 3]; + } + } else if (TK == 32) { + // the paired runs (element i + cap / 2 is (k + 16, n), the COOP_K4 layout): one W8 a pair + const uint16_t half_cap = rw.get_capacity() / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rw.get_multidimensional_index(i); + half4 w, w2; + W8::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab, w, w2); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rw.get_capacity(); i += 4) { + const auto ix = rw.get_multidimensional_index(i); + const half4 w = W4::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + } + tA = tH.template slice(k0, row0); + op.run(tA, rw, acc); + } + tensor, tensor_inline> tY(YE, dextents(a.n_embd, 1 << 30)); + acc.store(tY.slice(n0, row0)); +} + +#define MMGU3_(F, M, N, K, S) template [[host_name("mm_gu3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu3) mm_gu3; +#define MMDN3_(F, M, N, K, S) template [[host_name("mm_down3_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_down3) mm_down3; +#define MMGU3(F, M, N, K, S) MMGU3_(F, M, N, K, S) +#define MMDN3(F, M, N, K, S) MMDN3_(F, M, N, K, S) +#define MM3_ONE(F, M, N, K, S) MMGU3(F, M, N, K, S) MMDN3(F, M, N, K, S) +#define MM3_GU(F) MMGU3(F, 16, 64, 32, 2) MMGU3(F, 16, 32, 32, 4) MMGU3(F, 32, 64, 32, 2) MMGU3(F, 16, 64, 64, 2) \ + MMGU3(F, 16, 32, 16, 2) MMGU3(F, 32, 32, 32, 2) MMGU3(F, 16, 64, 32, 1) MMGU3(F, 16, 64, 32, 4) MMGU3(F, 16, 32, 16, 4) MMGU3(F, 16, 32, 32, 5) MMGU3(F, 16, 32, 16, 5) MMGU3(F, 16, 16, 32, 4) MMGU3(F, 16, 16, 32, 8) MMGU3(F, 16, 16, 16, 8) +#define MM3_DN(F) MMDN3(F, 16, 64, 32, 2) MMDN3(F, 16, 32, 32, 4) MMDN3(F, 32, 64, 32, 2) MMDN3(F, 16, 64, 64, 2) \ + MMDN3(F, 16, 32, 16, 2) MMDN3(F, 32, 32, 32, 2) MMDN3(F, 16, 64, 32, 1) MMDN3(F, 16, 64, 32, 4) MMDN3(F, 16, 64, 16, 2) MMDN3(F, 16, 32, 16, 4) MMDN3(F, 16, 64, 32, 5) MMDN3(F, 16, 64, 16, 4) +MM3_GU(T_Q4_K) MM3_GU(T_Q5_K) MM3_GU(T_Q6_K) MM3_GU(T_Q8_0) MM3_GU(T_IQ4_XS) MM3_GU(T_Q2_0) MM3_GU(T_IQ3_XXS) MM3_GU(T_IQ2_XXS) MM3_GU(T_IQ2_XS) MM3_GU(T_IQ2_S) MM3_GU(T_IQ3_S) +MM3_DN(T_Q2_0) MM3_DN(T_IQ4_NL) MM3_DN(T_Q5_1) MM3_DN(T_Q8_0) + +// mm_gu4: mm_gu3 with each decoded right operand (TN x TK weights of gate and of up) applied to MS 16-row token +// tiles (a tile of 16 * MS rows from the plan, tm = 16 * MS): the decode, which bounds mm_gu3, is shared by MS times +// as many token rows (more padding rows per expert, but the matrix ops are cheap next to the decode). The same +// per-row arithmetic as mm_gu3 (each 16-row op is mm_gu3's): bitwise for the rows both compute. COOP_K4 layout only. +template +kernel void mm_gu4(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16; + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* g0 = W + a.gate_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row; + device const uchar* u0 = W + a.up_off + ulong(e) * a.gu_stride + ulong(n0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rg = op.template get_right_input_cooperative_tensor(); + auto ru = op.template get_right_input_cooperative_tensor(); + auto tA0 = tX.template slice(0, row0); + auto tA1 = tX.template slice(0, row0 + TM); +#define MM4_DST op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>() + auto accg0 = MM4_DST; + auto accu0 = MM4_DST; + auto accg1 = MM4_DST; + auto accu1 = MM4_DST; +#undef MM4_DST + using DT = decltype(accg0); + for (uint16_t i = 0; i < accg0.get_capacity(); ++i) + if (accg0.is_valid_element(i)) { accg0[i] = 0.0f; accu0[i] = 0.0f; accg1[i] = 0.0f; accu1[i] = 0.0f; } + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + const uint16_t cap = rg.get_capacity(), half_cap = cap / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rg.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* rgw = g0 + ulong(n) * a.gu_row; + device const uchar* ruw = u0 + ulong(n) * a.gu_row; + half4 wg, wu, wg2, wu2; + W8::run(rgw, k, tab, wg, wg2); + W8::run(ruw, k, tab, wu, wu2); + rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; + ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; + ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + } + tA0 = tX.template slice(k0, row0); + op.run(tA0, rg, accg0); + op.run(tA0, ru, accu0); + if (MS > 1) { + tA1 = tX.template slice(k0, row0 + TM); + op.run(tA1, rg, accg1); + op.run(tA1, ru, accu1); + } + } + for (int s = 0; s < MS; ++s) { + thread DT& ag = s == 0 ? accg0 : accg1; + thread DT& au = s == 0 ? accu0 : accu1; + for (uint16_t i = 0; i < ag.get_capacity(); ++i) + if (ag.is_valid_element(i)) { + const float g = ag[i]; + ag[i] = g / (1.0f + exp(-g)) * au[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + simdgroup_barrier(mem_flags::mem_threadgroup); + ag.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + s * TM + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } + } +} +#define MMGU4_(F, M, N, K, S) template [[host_name("mm_gu4_" #F "_" #M "_" #N "_" #K "_" #S)]] kernel decltype(mm_gu4) mm_gu4; +#define MMGU4(F, M, N, K, S) MMGU4_(F, M, N, K, S) +#define MM4_GU(F) MMGU4(F, 2, 32, 32, 4) MMGU4(F, 2, 32, 32, 2) MMGU4(F, 2, 16, 32, 4) MMGU4(F, 1, 32, 32, 4) +MM4_GU(T_Q2_0) MM4_GU(T_IQ3_XXS) MM4_GU(T_IQ2_XXS) MM4_GU(T_IQ2_XS) MM4_GU(T_IQ2_S) MM4_GU(T_IQ3_S) MM4_GU(T_IQ4_XS) + +// mm_gu5: mm_gu3 (COOP_K4 layout, TK 32) with the threadgroup's weight rows staged block by block (QK = 256 values a +// row, gate and up) into threadgroup memory by coalesced ushort loads, the lanes decoding from there: the decode's +// many small scattered device loads (2-6 a run) become threadgroup reads. Same decoded values and the same matrix +// ops as mm_gu3: bitwise. Formats with 256-value blocks and a templated W8 (IQ3_S, IQ3_XXS, IQ2_S). +template +kernel void mm_gu5(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], + device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16, TK = 32, BB = fmt::BB, ROWS = SG * TN, RU = BB / 2; + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + threadgroup ushort stg[2][ROWS * RU]; + load_tab(tab, tid, SG * 32); + const int2 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, nt0 = int(tg.x) * ROWS, n0 = nt0 + sg * TN; + device const uchar* gt = W + a.gate_off + ulong(e) * a.gu_stride + ulong(nt0) * a.gu_row; + device const uchar* ut = W + a.up_off + ulong(e) * a.gu_stride + ulong(nt0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rg = op.template get_right_input_cooperative_tensor(); + auto ru = op.template get_right_input_cooperative_tensor(); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + const int NB = int(a.n_embd) / 256; + for (int kb = 0; kb < NB; ++kb) { + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int idx = tid; idx < ROWS * RU; idx += SG * 32) { + const int r = idx / RU, w = idx - r * RU; + stg[0][idx] = ((device const ushort*)(gt + ulong(r) * a.gu_row + kb * BB))[w]; + stg[1][idx] = ((device const ushort*)(ut + ulong(r) * a.gu_row + kb * BB))[w]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int ks = 0; ks < 8; ++ks) { + const int k0 = kb * 256 + ks * TK; + const uint16_t half_cap = rg.get_capacity() / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rg.get_multidimensional_index(i); + const int kl = ks * TK + ix[0], n = sg * TN + ix[1]; + // the block's row in the stage; W8 indexes block (k >> 8) = 0 of it with k = kl + threadgroup const uchar* rgw = (threadgroup const uchar*)(stg[0] + n * RU); + threadgroup const uchar* ruw = (threadgroup const uchar*)(stg[1] + n * RU); + half4 wg, wu, wg2, wu2; + W8::run(rgw, kl, tab, wg, wg2); + W8::run(ruw, kl, tab, wu, wu2); + rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; + ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; + ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + } + tA = tX.template slice(k0, row0); + op.run(tA, rg, accg); + op.run(tA, ru, accu); + } + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = accg[i]; + accg[i] = g / (1.0f + exp(-g)) * accu[i]; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + accg.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } +} +#define MMGU5_(F, N, S) template [[host_name("mm_gu5_" #F "_" #N "_" #S)]] kernel decltype(mm_gu5) mm_gu5; +#define MMGU5(F, N, S) MMGU5_(F, N, S) +MMGU5(T_IQ3_S, 32, 2) MMGU5(T_IQ3_S, 32, 1) MMGU5(T_IQ3_S, 16, 4) MMGU5(T_IQ3_XXS, 32, 2) MMGU5(T_IQ3_XXS, 32, 1) +MMGU5(T_IQ3_XXS, 16, 4) MMGU5(T_IQ2_S, 32, 2) MMGU5(T_IQ2_S, 32, 1) MMGU5(T_IQ2_S, 16, 4) diff --git a/edge0-strata/metal/kernels/prompt.metal b/edge0-strata/metal/kernels/prompt.metal new file mode 100644 index 00000000..64833b0a --- /dev/null +++ b/edge0-strata/metal/kernels/prompt.metal @@ -0,0 +1,177 @@ +// metal/kernels/prompt.metal - the elementwise steps of a prompt chunk around the GEMMs (after quant/units.metal): +// +// route_topk per token: softmax over the 512 logits, top 10 (ties: the lower expert), weights renormalised +// over the selected with a 2^-14 floor - moe_route's rule; one simdgroup a token +// gr_epi_down after the HC down GEMM (V = xn . [w_down; w_inject], T x (hc_lr + hc)): lo_h = half(silu(v / hc)), +// inject = v +// gr_epi_up after the HC up GEMM (G = lo . w_up, T x hc*n): mixed = mean_c xn[c] * sigmoid(G[c]) -> float and +// half (the next GEMMs' operand) +// swiglu_h h = half(silu(g) * u) +// moe_out bo = y_routed + sigmoid(sg) * y_shared + +kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]], + device float* w [[buffer(2)]], constant uint& n_tok [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const uint t = tg.x * 8 + sg; + if (t >= n_tok) return; + device const float* l = logits + ulong(t) * 512; + float v[16]; + for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; + float mx = -INFINITY; + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + if ((ex & 31) == lane) { + v[ex / 32] = -INFINITY; + ids[t * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + if (lane < 10) w[t * 10 + lane] = selected * (1.0f / ssum); +} + +struct EpiArgs { + uint n_tok; + uint n_embd; // 2560 + uint hc; // 4 + uint hc_lr; // 320 +}; + +// grid (T * (hc_lr + hc) / 256) +kernel void gr_epi_down(device const float* V [[buffer(0)]], device half* lo_h [[buffer(1)]], + device float* inject [[buffer(2)]], constant EpiArgs& a [[buffer(3)]], + uint gid [[thread_position_in_grid]]) { + const uint nr = a.hc_lr + a.hc, t = gid / nr, r = gid % nr; + if (t >= a.n_tok) return; + const float v = V[gid]; + if (r < a.hc_lr) { + const float z = v / float(a.hc); + lo_h[t * a.hc_lr + r] = half(z / (1.0f + exp(-z))); + } else { + inject[t * a.hc + r - a.hc_lr] = v; + } +} + +// grid (T * n_embd / 256) +kernel void gr_epi_up(device const float* G [[buffer(0)]], device const float* xn [[buffer(1)]], + device float* mixed [[buffer(2)]], device half* mixed_h [[buffer(3)]], + constant EpiArgs& a [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + const uint n = a.n_embd, t = gid / n, d = gid % n; + if (t >= a.n_tok) return; + float m = 0.0f; + for (uint c = 0; c < a.hc; ++c) { + const float g = G[ulong(t) * a.hc * n + c * n + d]; + m = fma(xn[ulong(t) * a.hc * n + c * n + d], 1.0f / (1.0f + exp(-g)), m); + } + m = m / float(a.hc); + mixed[gid] = m; + mixed_h[gid] = half(m); +} + +kernel void swiglu_h(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], device half* h [[buffer(2)]], + constant uint& n [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + if (gid >= n) return; + const float gv = g[gid]; + h[gid] = half(gv / (1.0f + exp(-gv)) * u[gid]); +} + +// grid (T * n_embd / 256); sg: T values +kernel void moe_out(device const float* yr [[buffer(0)]], device const float* ys [[buffer(1)]], + device const float* sg [[buffer(2)]], device float* bo [[buffer(3)]], constant uint2& a [[buffer(4)]], + uint gid [[thread_position_in_grid]]) { + const uint t = gid / a.y; + if (t >= a.x) return; + bo[gid] = fma(1.0f / (1.0f + exp(-sg[t])), ys[gid], yr[gid]); +} + +// ---- the HC read with fewer passes (after dense_mm.metal: GemmArgs, MetalPerformancePrimitives) --------------------- +// gr_norm_h: gr_norm with xn written as fp16 only (the GEMM operand; gemm_hc_up's epilogue reads it too). +// one threadgroup per (stream c, token t), 256 threads; args as gr_norm's (GrArgs, gr.metal) +kernel void gr_norm_h(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]], + device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]], + device half* xnh [[buffer(4)]], constant GrArgs& a [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device float* r = R + (t * a.hc + c) * n; + const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f; + float ss = 0.0f; + for (uint d = tid; d < n; d += 256) { + float v = r[d]; + if (a.apply) { + v = fma(bo[t * n + d], g, v); + r[d] = v; + } + ss = fma(v, v, ss); + } + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(n) + a.eps); + device half* o = xnh + (t * a.hc + c) * n; + device const float* w = w_norm + c * n; + for (uint d = tid; d < n; d += 256) o[d] = half(r[d] * w[d] * rs); +} + +// gemm_hc_up: G = lo . w_up^T with w_up's rows interleaved by dq_f16 (perm = hc: row d * hc + c), so a TN-column +// tile holds TN / hc outputs d with all their streams; the epilogue mixed[t][d] = mean_c xn[t][c][d] * +// sigmoid(G[t][d * hc + c]) -> fp32 and fp16. G is never written. grid (hc * n / TN, ceil(T / TM)), 32 * SG threads. +template +kernel void gemm_hc_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const half* xnh [[buffer(2)]], + device float* mixed [[buffer(3)]], device half* mixh [[buffer(4)]], constant GemmArgs& g [[buffer(5)]], + constant EpiArgs& e [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float ot[TM * TN]; + tensor, tensor_inline> tA(A, dextents(g.K, g.M)); + tensor, tensor_inline> tB(B, dextents(g.K, g.N)); + tensor, tensor_inline> tO(ot, dextents(TN, TM)); + constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tA.slice(0, int(tg.y) * TM); + auto mB = tB.slice(0, int(tg.x) * TN); + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int hc = int(e.hc), n = int(e.n_embd), DPT = TN / hc; + const int t0 = int(tg.y) * TM, d0 = int(tg.x) * DPT; + for (int idx = tid; idx < TM * DPT; idx += SG * 32) { + const int m = idx / DPT, dd = idx - m * DPT, t = t0 + m; + if (t >= int(e.n_tok)) continue; + float s = 0.0f; + for (int c = 0; c < hc; ++c) + s = fma(float(xnh[ulong(t) * hc * n + c * n + d0 + dd]), 1.0f / (1.0f + exp(-ot[m * TN + dd * hc + c])), s); + s = s / float(hc); + mixed[ulong(t) * n + d0 + dd] = s; + mixh[ulong(t) * n + d0 + dd] = half(s); + } +} +template [[host_name("gemm_hc_up_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4>) gemm_hc_up<64, 64, 4>; +template [[host_name("gemm_hc_up_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4>) gemm_hc_up<64, 32, 4>; +template [[host_name("gemm_hc_up_32_64_4")]] kernel decltype(gemm_hc_up<32, 64, 4>) gemm_hc_up<32, 64, 4>; diff --git a/edge0-strata/metal/kernels/qsa.metal b/edge0-strata/metal/kernels/qsa.metal new file mode 100644 index 00000000..1cba1bb8 --- /dev/null +++ b/edge0-strata/metal/kernels/qsa.metal @@ -0,0 +1,409 @@ +// metal/kernels/qsa.metal - QSA (sparse full attention) decode for a window of tokens, between the projections. +// Contract of include/strata/kernels/qsa.hpp / qsa_select.hpp / qsa_decode_attn.hpp / kv_q8.hpp (the engine's +// native path), in a Metal layout (no paging: unified memory): +// +// qsa_prep per token t (cell pos0 + t): q heads rms_norm * q_norm, NEOX RoPE on 64 dims -> qcur; +// k rms_norm * k_norm, RoPE -> int8 KV append (fp16 scale per 64 = fp16(amax/127), rint, clamp); +// v -> int8 append; indexer q rms_norm * iq_norm, RoPE -> qidx. The indexer KEYS (last +// threadgroup, all window tokens in order): raw tail [cell % 4]; on a block completion +// pooled[b] = rope(rms_norm(mean of the block's 4 raw keys) * ik_norm, pos of the block's first cell); +// cell 0 also sets dead = rms_norm(raw) * ik_norm (the tail block's key). +// qsa_score per token, block b <= n_bid: sum over the 4 indexer heads of relu(qidx . key_b); key_b = pooled[b], +// the tail block n_bid scores `dead` and gets +1e9 when it has cells. +// qsa_topk per token: width = min(n_kv, 2051) cells; n_kv <= 2051: all; else the width cells of largest +// score (a cell inherits its block's), ties to the lowest cell, emitted ascending (radix select). +// qsa_attn split-K: (chunk of CH cells, kv head, token): 12 query heads; s = q.k / 16, chunk max m, sum l, +// acc = p . v -> partials +// qsa_merge per (head, token): log-sum-exp merge, * sigmoid(gate), written as prepared activations (xprep +// planes) for attn_output. +// +// RoPE: cos/sin from a host table [pos][32] (double precision, as build_rope_table). + +#define QH 24 +#define QKV 2 +#define QD 256 +#define QROT 64 +#define IH 4 +#define ID 128 +#define QBUDGET 2051 +#define QCH 64 + +constant int QSA_NTOK [[function_constant(3)]]; +// QSA_MROPE (function constant 24): interleaved M-RoPE (the vision path; strata/kernels/mrope.hpp, ggml IMROPE, +// sections 11/11/10): buffer 20 of qsa_prep is a table int32 [cells][3] (t, h, w), and rotary pair i (0..31) of the row +// at (pos_base + cell) rotates by the cos/sin table's row tab[(pos_base + cell) * 3 + i % 3] instead of its own - +// the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE. +constant bool QSA_MROPE [[function_constant(24)]]; +constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE; +constant bool QSA_ROWS_C [[function_constant(30)]]; +constant bool qsa_rows = is_function_constant_defined(QSA_ROWS_C) && QSA_ROWS_C; +// QSA_ROWS (function constant 30): the rows of the window are DIFFERENT sequences (batch slots), one token each, each at its own +// cell and in its own slot of the caches. A buffer rinfo (uint) carries the strides between the slots' caches, in elements: +// [0] kq / vq (chars), [1] ksc / vsc (halves), [2] tail, [3] dead, [4] pooled (floats), [5] GDN state (floats), [6] conv history +// (floats, per slot and parity), [7] PLE history (floats); then row t at [8 + 4 t] = {slot, cell, history parity, 0}. Off: exactly the +// plain kernels. + +struct QsaArgs { + uint pos0; // the window's first cell (relative to the sequence) + uint pos_base; // absolute position of cell 0 (RoPE) + uint max_cells; + uint n_tok; + float eps; +}; + +// cache layout per layer +// kq/vq [cell][kv head][256] int8, ks/vs [cell][kv head][4] half +// tail [4][128] f32 (slot = cell % 4), dead [128], pooled [max_cells / 4 + 1][128] + +static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushort lane, ushort nsg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; + for (int i = 0; i < nsg; ++i) t += part[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return t; +} + +// one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32) +static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row, + threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps, + device const float2* cs, uint absp, bool rope, device const int* mtab) { + const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8); + const float rs = rsqrt(ss / float(w) + eps); + const float y = tid < w ? (rs * x) * g[tid] : 0.0f; + row[tid] = y; + threadgroup_barrier(mem_flags::mem_threadgroup); + float o = y; + if (rope && tid < QROT) { + const int p = tid & 31; + const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp; + const float2 c = cs[pp * 32 + p]; + const float a = row[p], b = row[p + 32]; + o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + (void)xhi_unused; + return o; +} + +// grid (32 rows + 1, n_tok): rows 0..23 q heads, 24..25 k heads, 26..27 v heads, 28..31 indexer q heads; +// threadgroup 32 (x = 32, y = 0) is the indexer-key threadgroup. 256 threads. +kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float* kcur [[buffer(1)]], + device const float* vcur [[buffer(2)]], device const float* iqraw [[buffer(3)]], + device const float* ikraw [[buffer(4)]], device const float* qn [[buffer(5)]], + device const float* kn [[buffer(6)]], device const float* iqn [[buffer(7)]], + device const float* ikn [[buffer(8)]], device const float2* cs [[buffer(9)]], + device float* qcur [[buffer(10)]], device float* qidx [[buffer(11)]], + device char* kq [[buffer(12)]], device char* vq [[buffer(13)]], device half* ksc [[buffer(14)]], + device half* vsc [[buffer(15)]], device float* tail [[buffer(16)]], device float* dead [[buffer(17)]], + device float* pooled [[buffer(18)]], constant QsaArgs& a [[buffer(19)]], + device const int* mtab [[buffer(20)]], + device const uint* rinfo [[buffer(21), function_constant(qsa_rows)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + threadgroup float row[256]; + threadgroup float amax[8]; + const int r = int(tg.x); + if (r == 32) { + // the indexer keys, every window token in order (only threadgroup (32, 0) works) + if (tg.y != 0) return; + for (int t = 0; t < int(a.n_tok); ++t) { + uint cell = a.pos0 + uint(t); + device float* tail_t = tail; device float* dead_t = dead; device float* pooled_t = pooled; + if (qsa_rows) { + const ulong sl = rinfo[8 + 4 * t]; + cell = rinfo[8 + 4 * t + 1]; + tail_t = tail + sl * rinfo[2]; dead_t = dead + sl * rinfo[3]; pooled_t = pooled + sl * rinfo[4]; + } + const float x = tid < ID ? ikraw[t * ID + tid] : 0.0f; + if (tid < ID) tail_t[(cell & 3) * ID + tid] = x; + threadgroup_barrier(mem_flags::mem_device); + if (cell == 0) { + const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab); + if (tid < ID) dead_t[tid] = o; + } + if ((cell & 3) == 3) { + float m = 0.0f; + if (tid < ID) m = (((tail_t[tid] + tail_t[ID + tid]) + tail_t[2 * ID + tid]) + tail_t[3 * ID + tid]) * 0.25f; + const uint b = cell >> 2; + const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab); + if (tid < ID) pooled_t[b * ID + tid] = o; + } + threadgroup_barrier(mem_flags::mem_device); + } + return; + } + const int t = int(tg.y); + const uint cell = qsa_rows ? rinfo[8 + 4 * t + 1] : a.pos0 + uint(t), absp = a.pos_base + cell; + device char* kqr = kq; device char* vqr = vq; device half* kscr = ksc; device half* vscr = vsc; + if (qsa_rows) { + const ulong sl = rinfo[8 + 4 * t]; + kqr = kq + sl * rinfo[0]; vqr = vq + sl * rinfo[0]; kscr = ksc + sl * rinfo[1]; vscr = vsc + sl * rinfo[1]; + } + if (r < QH) { + const float x = qfull[(t * QH + r) * 2 * QD + tid]; + const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + qcur[(t * QH + r) * QD + tid] = o; + } else if (r < QH + 2 * QKV) { + const bool isv = r >= QH + QKV; + const int h = r - QH - (isv ? QKV : 0); + float x = (isv ? vcur : kcur)[(t * QKV + h) * QD + tid]; + if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab); + // int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads) + float m = simd_max(fabs(x)); + if (lane == 0) amax[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + const int g = tid >> 6; + m = max(amax[2 * g], amax[2 * g + 1]); + const half sh = half(m / 127.0f); + const float sf = float(sh); + int q = 0; + if (sf > 0.0f) q = clamp(int(rint(x / sf)), -127, 127); + const uint rowi = cell * QKV + uint(h); + (isv ? vqr : kqr)[rowi * QD + tid] = char(q); + if ((tid & 63) == 0) (isv ? vscr : kscr)[rowi * 4 + g] = sh; + } else { + const int h = r - QH - 2 * QKV; + const float x = tid < ID ? iqraw[(t * IH + h) * ID + tid] : 0.0f; + const float o = norm_rope(x, 0.0f, iqn, row, part, tid, sg, lane, ID, a.eps, cs, absp, true, mtab); + if (tid < ID) qidx[(t * IH + h) * ID + tid] = o; + } +} + +// scores [n_tok][max_blocks]; grid (ceil(max_blocks / 8), n_tok), 256 threads: a simdgroup per block. +kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + device const uint* rinfo [[buffer(5), function_constant(qsa_rows)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const int t = int(tg.y); + const uint n_kv = (qsa_rows ? rinfo[8 + 4 * t + 1] : a.pos0 + uint(t)) + 1, n_bid = n_kv >> 2; + const uint b = tg.x * 8 + sg; + if (b > n_bid) return; + device const float* pooled_s = pooled; device const float* dead_s = dead; + if (qsa_rows) { const ulong sl = rinfo[8 + 4 * t]; pooled_s = pooled + sl * rinfo[4]; dead_s = dead + sl * rinfo[3]; } + device const float* key = b == n_bid ? dead_s : pooled_s + b * ID; + float k[4]; + for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i]; + float s = 0.0f; + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + for (int i = 0; i < 4; ++i) d = fma(qidx[(t * IH + h) * ID + lane + 32 * i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if (b == n_bid && (n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * (a.max_cells / 4 + 2) + b] = s; +} + +// ids [n_tok][QBUDGET], ascending. One threadgroup of 1024 threads per token. Order key of a float: larger score = +// larger key; a radix select (4 passes of 8 bits) over the blocks with weights = their cell counts finds the +// threshold key T and how many cells with key == T to take (the lowest cells first). +static inline uint okey(float f) { + const uint u = as_type(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +kernel void qsa_topk(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sel[4]; + threadgroup uint pref[32]; + const int t = int(tg.x); + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + // radix select of the `need`-th largest cell key + uint prefix = 0, pmask = 0, need = QBUDGET; + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint b = tid; b < nb; b += 1024) { + const uint k = okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + uint acc = 0, d = 0; + for (int v = 255; v >= 0; --v) { + const uint h = atomic_load_explicit(&hist[v], memory_order_relaxed); + if (acc + h >= need) { d = uint(v); break; } + acc += h; + } + sel[0] = d; + sel[1] = need - acc; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + // T = prefix; take every cell with key > T and the first `need` cells (lowest index) with key == T, ascending. + // Pass over blocks in order of index with a prefix scan of each block's selected cell count. + const uint T = prefix; + uint base = 0, eq_taken = 0; + for (uint b0 = 0; b0 < nb; b0 += 1024) { + const uint b = b0 + tid; + uint w = 0, weq = 0; + if (b < nb) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) w = cells; + else if (k == T) weq = cells; + } + // equal-key cells: the lowest first, `need` of them in total + const uint eq_before = simd_prefix_exclusive_sum(weq); + if (lane == 31) pref[sg] = eq_before + weq; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint eq_sg = 0; + for (int i = 0; i < int(sg); ++i) eq_sg += pref[i]; + uint eq_tot = 0; + for (int i = 0; i < 32; ++i) eq_tot += pref[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint eq0 = eq_taken + eq_sg + eq_before; + uint take_eq = 0; + if (weq && eq0 < need) take_eq = min(weq, need - eq0); + const uint mine = w + take_eq; + const uint before = simd_prefix_exclusive_sum(mine); + if (lane == 31) pref[sg] = before + mine; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint off = base; + for (int i = 0; i < int(sg); ++i) off += pref[i]; + uint tot = 0; + for (int i = 0; i < 32; ++i) tot += pref[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + off += before; + for (uint c = 0; c < mine; ++c) out[off + c] = int(4 * b + c); + base += tot; + eq_taken += eq_tot; + } +} + +// partials: [n_tok][chunks][QH][QD + 2] (acc, then m, l). grid (chunks, QKV, n_tok), 256 threads. +kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + device const uint* rinfo [[buffer(8), function_constant(qsa_rows)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int HPG = QH / QKV; // 12 query heads per kv head + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = (qsa_rows ? rinfo[8 + 4 * t + 1] : a.pos0 + uint(t)) + 1, width = min(n_kv, uint(QBUDGET)); + if (qsa_rows) { const ulong sl = rinfo[8 + 4 * t]; kq += sl * rinfo[0]; vq += sl * rinfo[0]; ksc += sl * rinfo[1]; vsc += sl * rinfo[1]; } + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) { + // an empty chunk: m = -inf, l = 0 (the merge skips it) + for (int h = 0; h < HPG; ++h) + if (tid == 0) { pp[(g * HPG + h) * (QD + 2) + QD] = -INFINITY; pp[(g * HPG + h) * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + // scores: 12 x 64 dots of 256; simdgroup sg takes cells sg, sg + 8, ...; lanes split the 256 dims (8 each) + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char* kr = kq + rowi * QD; + float kv[8]; + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + for (int i = 0; i < 8; ++i) kv[i] = float(kr[lane * 8 + i]) * scl; + for (int h = 0; h < HPG; ++h) { + float d = 0.0f; + for (int i = 0; i < 8; ++i) d = fma(qs[h][lane * 8 + i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // per head: max, exp, sum (simdgroup h % 8 handles heads h, h + 8) + threadgroup float mh[HPG], lh[HPG]; + for (int h = sg; h < HPG; h += 8) { + float m = -INFINITY; + for (int c = lane; c < nc; c += 32) m = max(m, s[h][c]); + m = simd_max(m); + float l = 0.0f; + for (int c = lane; c < nc; c += 32) { + const float p = exp(s[h][c] - m); + s[h][c] = p; + l += p; + } + l = simd_sum(l); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // acc[h][j] = sum_c p[h][c] v[c][j]: thread = dim j + float acc[HPG]; + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + for (int c = 0; c < nc; ++c) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + const float v = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][c], v, acc[h]); + } + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +// grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token +kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + device const uint* rinfo [[buffer(4), function_constant(qsa_rows)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float msum[256]; + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH - 1) / QCH; + const uint n_kv = (qsa_rows ? rinfo[8 + 4 * t + 1] : a.pos0 + uint(t)) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH - 1) / QCH); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + float M = -INFINITY; + for (int k = 0; k < nc; ++k) M = max(M, pp[k * QH * (QD + 2) + QD]); + float L = 0.0f, o = 0.0f; + for (int k = 0; k < nc; ++k) { + device const float* q = pp + k * QH * (QD + 2); + const float e = exp(q[QD] - M); + L = fma(q[QD + 1], e, L); + o = fma(q[tid], e, o); + } + o = o / L; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + // prepared planes (xprep layout) of the n = QH * QD vector + const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; + const int j = tid & 15; + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; + const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + device float* raw = xp + uint(t) * n + d; + raw[0] = y; + raw[plane] = y * q2s[j & 7]; + raw[2 * plane] = y * p4s[j & 3]; + raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; + msum[tid] = y; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (j == 0) { + float s0 = 0.0f, s1 = 0.0f; + for (int k = 0; k < 16; k += 2) { s0 += msum[tid + k]; s1 += msum[tid + k + 1]; } + xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; + } +} diff --git a/edge0-strata/metal/kernels/qsa2.metal b/edge0-strata/metal/kernels/qsa2.metal new file mode 100644 index 00000000..94ca9a38 --- /dev/null +++ b/edge0-strata/metal/kernels/qsa2.metal @@ -0,0 +1,290 @@ +// metal/kernels/qsa2.metal - QSA decode v2 (after qsa.metal): the same selection and attention math, more parallel. +// +// qsa_attn2 chunks of 32 cells, the 12 query heads of a kv head in 2 groups of 6: 4x the threadgroups of +// qsa_attn (at a 1K context 128 instead of 32); V loads 8 cells at a time. +// qsa_merge2 qsa_merge for that chunking. +// qsa_topk2 the digit of each radix pass found by a parallel suffix scan of the 256 bins (simdgroup prefix +// sums) instead of one thread; cells emitted from contiguous block ranges per thread after one scan. + +#define QCH2 32 +#define QHG 2 + +// partials: [n_tok][chunks2][QH][QD + 2]. grid (chunks2, QKV * QHG, n_tok), 256 threads. +kernel void qsa_attn2(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int HPG = QH / QKV / QHG; // 6 + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH2]; + threadgroup float mh[HPG], lh[HPG]; + threadgroup int cid[QCH2]; + const int chunk = int(tg.x), g = int(tg.y) / QHG, hg = int(tg.y) % QHG, t = int(tg.z); + const int h0 = g * (QH / QKV) + hg * HPG; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH2; + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + if (uint(c0) >= width) return; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + const int nc = min(QCH2, int(width) - c0); + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + h0 + h) * QD + tid]; + if (tid < QCH2) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int c = sg; c < nc; c += 8) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD + lane * 8); + const char4 k0 = kr[0], k1 = kr[1]; + const float scl = float(ksc[rowi * 4 + (lane >> 3)]); + const float kv[8] = {float(k0.x) * scl, float(k0.y) * scl, float(k0.z) * scl, float(k0.w) * scl, + float(k1.x) * scl, float(k1.y) * scl, float(k1.z) * scl, float(k1.w) * scl}; + for (int h = 0; h < HPG; ++h) { + float d = 0.0f; + for (int i = 0; i < 8; ++i) d = fma(qs[h][lane * 8 + i], kv[i], d); + d = simd_sum(d); + if (lane == 0) s[h][c] = d * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg < HPG) { + const int h = sg; + const float sv = lane < nc ? s[h][lane] : -INFINITY; + const float m = simd_max(sv); + const float p = lane < nc ? exp(sv - m) : 0.0f; + const float l = simd_sum(p); + if (lane < nc) s[h][lane] = p; + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + for (int cb = 0; cb < nc; cb += 8) { + float v[8]; + const int m = min(8, nc - cb); + for (int c = 0; c < 8; ++c) { + v[c] = 0.0f; + if (c < m) { + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + } + } + for (int c = 0; c < m; ++c) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (h0 + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} + +kernel void qsa_merge2(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]], + device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float msum[256]; + const int h = int(tg.x), t = int(tg.y); + const int nchunks = (QBUDGET + QCH2 - 1) / QCH2; + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nc = int((width + QCH2 - 1) / QCH2); + device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2); + float M = -INFINITY; + for (int k = 0; k < nc; ++k) M = max(M, pp[k * QH * (QD + 2) + QD]); + float L = 0.0f, o = 0.0f; + for (int k = 0; k < nc; ++k) { + device const float* q = pp + k * QH * (QD + 2); + const float e = exp(q[QD] - M); + L = fma(q[QD + 1], e, L); + o = fma(q[tid], e, o); + } + o = o / L; + const float gate = qfull[(t * QH + h) * 2 * QD + QD + tid]; + const float y = o * (1.0f / (1.0f + exp(-gate))); + const uint n = QH * QD, plane = a.n_tok * n, d = uint(h) * QD + tid; + const int j = tid & 15; + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; + const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + device float* raw = xp + uint(t) * n + d; + raw[0] = y; + raw[plane] = y * q2s[j & 7]; + raw[2 * plane] = y * p4s[j & 3]; + raw[3 * plane] = (j & 1) ? y * 0.00390625f : y; + msum[tid] = y; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (j == 0) { + float s0 = 0.0f, s1 = 0.0f; + for (int k = 0; k < 16; k += 2) { s0 += msum[tid + k]; s1 += msum[tid + k + 1]; } + xp[4 * plane + uint(t) * (n / 16) + d / 16] = s0 + s1; + } +} + +kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]], + constant QsaArgs& a [[buffer(2)]], + device const uint* rinfo [[buffer(3), function_constant(qsa_rows)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup atomic_uint hist[256]; + threadgroup uint sgs[32], sgs2[32]; + threadgroup uint sel[2]; + const int t = int(tg.x); + const uint n_kv = (qsa_rows ? rinfo[8 + 4 * t + 1] : a.pos0 + uint(t)) + 1, n_bid = n_kv >> 2, nb = n_bid + 1; + device int* out = ids + t * QBUDGET; + device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2); + if (n_kv <= QBUDGET) { + for (uint c = tid; c < n_kv; c += 1024) out[c] = int(c); + return; + } + uint prefix = 0, pmask = 0, need = QBUDGET; + for (int pass = 0; pass < 4; ++pass) { + const int shift = 24 - 8 * pass; + if (tid < 256) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint b = tid; b < nb; b += 1024) { + const uint k = okey(sc[b]); + if ((k & pmask) == prefix) { + const uint w = b < n_bid ? 4u : (n_kv & 3u); + if (w) atomic_fetch_add_explicit(&hist[(k >> shift) & 255u], w, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // thread i < 256 holds bin 255 - i: an exclusive prefix over i = the cells of the higher digits + uint cnt = 0, ex = 0; + if (tid < 256) { + cnt = atomic_load_explicit(&hist[255 - tid], memory_order_relaxed); + ex = simd_prefix_exclusive_sum(cnt); + if (lane == 31) sgs[sg] = ex + cnt; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 256) { + uint above = ex; + for (int i = 0; i < int(sg); ++i) above += sgs[i]; + if (above < need && above + cnt >= need) { sel[0] = 255u - tid; sel[1] = need - above; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + prefix |= sel[0] << shift; + pmask |= 255u << shift; + need = sel[1]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const uint T = prefix; + const uint per = (nb + 1023) / 1024, b0 = min(nb, uint(tid) * per), b1 = min(nb, b0 + per); + uint gt = 0, eq = 0; + for (uint b = b0; b < b1; ++b) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + if (k > T) gt += cells; + else if (k == T) eq += cells; + } + const uint eqb = simd_prefix_exclusive_sum(eq); + if (lane == 31) sgs2[sg] = eqb + eq; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint eq_before = eqb; + for (int i = 0; i < int(sg); ++i) eq_before += sgs2[i]; + const uint eq_take = eq_before >= need ? 0u : min(eq, need - eq_before); + const uint mine = gt + eq_take; + const uint mb = simd_prefix_exclusive_sum(mine); + if (lane == 31) sgs[sg] = mb + mine; + threadgroup_barrier(mem_flags::mem_threadgroup); + uint off = mb; + for (int i = 0; i < int(sg); ++i) off += sgs[i]; + uint left = eq_take; + for (uint b = b0; b < b1; ++b) { + const uint k = okey(sc[b]); + const uint cells = b < n_bid ? 4u : (n_kv & 3u); + uint take = 0; + if (k > T) take = cells; + else if (k == T) { take = min(cells, left); left -= take; } + for (uint c = 0; c < take; ++c) out[off + c] = int(4 * b + c); + off += take; + } +} + +// qsa_attn3: qsa_attn's chunking and partial layout (qsa_merge reads it), without a reduction per (cell, head): +// for q.k a lane owns a cell (8 simdgroups = 2 halves of the 64 cells x 4 groups of 3 heads; q from threadgroup +// memory, read by every lane alike), for p.v a thread owns a dimension. grid (33, QKV, n_tok), 256 threads. +kernel void qsa_attn3(device const float* qcur [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int HPG = QH / QKV; // 12 + threadgroup float qs[HPG][QD]; + threadgroup float s[HPG][QCH]; + threadgroup float mh[HPG], lh[HPG]; + threadgroup int cid[QCH]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + const int nchunks = (QBUDGET + QCH - 1) / QCH; + device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2); + if (uint(c0) >= width) return; + const int nc = min(QCH, int(width) - c0); + for (int h = 0; h < HPG; ++h) qs[h][tid] = qcur[(t * QH + g * HPG + h) * QD + tid]; + if (tid < QCH) cid[tid] = tid < nc ? ids[t * QBUDGET + c0 + tid] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + { + const int half_ = sg & 1, hgrp = sg >> 1; // cells 32 * half_ + lane, heads 3 * hgrp .. + 3 + const int c = 32 * half_ + lane; + if (c < nc) { + const uint rowi = uint(cid[c]) * QKV + uint(g); + device const char4* kr = (device const char4*)(kq + rowi * QD); + float d0 = 0.0f, d1 = 0.0f, d2 = 0.0f; + for (int gq = 0; gq < 4; ++gq) { + const float scl = float(ksc[rowi * 4 + gq]); + float e0 = 0.0f, e1 = 0.0f, e2 = 0.0f; + for (int i = 0; i < 16; ++i) { + const char4 k4 = kr[gq * 16 + i]; + const int base = gq * 64 + 4 * i; + const float kv[4] = {float(k4.x), float(k4.y), float(k4.z), float(k4.w)}; + for (int e = 0; e < 4; ++e) { + e0 = fma(qs[3 * hgrp][base + e], kv[e], e0); + e1 = fma(qs[3 * hgrp + 1][base + e], kv[e], e1); + e2 = fma(qs[3 * hgrp + 2][base + e], kv[e], e2); + } + } + d0 = fma(e0, scl, d0); + d1 = fma(e1, scl, d1); + d2 = fma(e2, scl, d2); + } + s[3 * hgrp][c] = d0 * 0.0625f; + s[3 * hgrp + 1][c] = d1 * 0.0625f; + s[3 * hgrp + 2][c] = d2 * 0.0625f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int h = sg; h < HPG; h += 8) { + float m = -INFINITY; + for (int c = lane; c < nc; c += 32) m = max(m, s[h][c]); + m = simd_max(m); + float l = 0.0f; + for (int c = lane; c < nc; c += 32) { + const float p = exp(s[h][c] - m); + s[h][c] = p; + l += p; + } + l = simd_sum(l); + if (lane == 0) { mh[h] = m; lh[h] = l; } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float acc[HPG]; + for (int h = 0; h < HPG; ++h) acc[h] = 0.0f; + for (int cb = 0; cb < nc; cb += 8) { + float v[8]; + const int m = min(8, nc - cb); + for (int c = 0; c < 8; ++c) { + v[c] = 0.0f; + if (c < m) { + const uint rowi = uint(cid[cb + c]) * QKV + uint(g); + v[c] = float(vq[rowi * QD + tid]) * float(vsc[rowi * 4 + (tid >> 6)]); + } + } + for (int c = 0; c < m; ++c) + for (int h = 0; h < HPG; ++h) acc[h] = fma(s[h][cb + c], v[c], acc[h]); + } + for (int h = 0; h < HPG; ++h) { + device float* o = pp + (g * HPG + h) * (QD + 2); + o[tid] = acc[h]; + if (tid == 0) { o[QD] = mh[h]; o[QD + 1] = lh[h]; } + } +} diff --git a/edge0-strata/metal/kernels/qsa_prompt.metal b/edge0-strata/metal/kernels/qsa_prompt.metal new file mode 100644 index 00000000..7029684e --- /dev/null +++ b/edge0-strata/metal/kernels/qsa_prompt.metal @@ -0,0 +1,348 @@ +// metal/kernels/qsa_prompt.metal - QSA attention for a PROMPT chunk on the matrix units (after qsa.metal: QH, QKV, QD, +// QBUDGET, QCH, QsaArgs). The same split-K partials as qsa_attn (part [t][chunk][QH][QD + 2]: acc, then m, l), so +// qsa_merge is unchanged; per (chunk of QCH = 64 selected cells, kv head g, token t) one simdgroup: +// S[12 x 64] = Qh[12 x 256] . K^T / 16 single-simdgroup matmul2d, K (int8 * fp16 scale per 64) decoded into +// the right-input cooperative tensor, 32 dims a step +// per head: m = max, p = exp(s - m), l = sum p (fp32, through threadgroup memory), P -> fp16 +// acc[12 x 256] = P[12 x 64] . V[64 x 256] four 64-dim matmuls, V decoded into the right input +// The operands are fp16 (q rounded from fp32; k, v = int8 * scale rounded): to rounding of qsa_attn's fp32. +#include +using namespace mpp::tensor_ops; +#ifndef STRATA_COOP_K4 +#define STRATA_COOP_K4 +// function constant 7: the right-input cooperative tensor holds runs of 4 consecutive k at one n per lane (the M5's +// layout; tests/coop_layout). false (the M4: runs of 2 n at one k): every element decoded on its own. +constant bool COOP_K4 [[function_constant(7)]]; +constant bool coop_k4 = is_function_constant_defined(COOP_K4) ? COOP_K4 : true; +#endif + +// a bisection probe (function constant 8): stop at stage n and write m = n (1: after the ids, 2: after S, 3: after +// the softmax); 0 = normal +constant int QSA_STOP [[function_constant(8)]]; +constant int qsa_stop = is_function_constant_defined(QSA_STOP) ? QSA_STOP : 0; + +// grid (chunks, QKV, n_tok), 32 threads. qh: the q heads as fp16 [t][QH][QD] (qcur rounded). +kernel void qsa_attn_mm(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]], + uint3 tg [[threadgroup_position_in_grid]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int HPG = QH / QKV, NCH = (QBUDGET + QCH - 1) / QCH; + threadgroup float s[16 * QCH]; + threadgroup half p[16 * QCH]; + threadgroup int cid[QCH]; + threadgroup half vs[QCH][4]; + const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int c0 = chunk * QCH; + device float* pp = part + ((uint(t) * NCH + chunk) * QH + g * HPG) * (QD + 2); + if (uint(c0) >= width) { + if (lane < HPG) { pp[lane * (QD + 2) + QD] = -INFINITY; pp[lane * (QD + 2) + QD + 1] = 0.0f; } + return; + } + const int nc = min(QCH, int(width) - c0); + for (int c = lane; c < QCH; c += 32) { + const int id = c < nc ? ids[t * QBUDGET + c0 + c] : -1; + cid[c] = id; + for (int j = 0; j < 4; ++j) vs[c][j] = id >= 0 ? vsc[(uint(id) * QKV + g) * 4 + j] : half(0.0f); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 1) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 1.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- S = Q K^T + { + constexpr auto desc = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD, + dextents(QD, HPG)); + auto rk = op.template get_right_input_cooperative_tensor(); + auto tA = tQ.template slice<32, 16>(0, 0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int k0 = 0; k0 < QD; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + // the row is clamped and the select is on the value: a select with the device load in the arm + // loses the threadgroup on the M1 (macOS 26, measured); the discarded value cannot matter + const uint row = uint(max(id, 0)) * QKV + g; + const float v = float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)]); + rk[i] = id >= 0 ? half(v) : half(0.0f); + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], c = ix[1], id = cid[c]; + half4 w = half4(0.0h); + if (id >= 0) { + const uint row = uint(id) * QKV + g; + const char4 q4 = *(device const char4*)(kq + row * QD + k); + w = half4(float4(q4) * float(ksc[row * 4 + (k >> 6)])); + } + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + op.run(tA, rk, acc); + } + tensor, tensor_inline> tS(s, dextents(QCH, 16)); + acc.store(tS); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 2) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 2.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- softmax pieces per head (lane = head for the reductions; 12 heads) + float mh = -INFINITY, lh = 0.0f; + if (lane < HPG) { + for (int c = 0; c < nc; ++c) mh = max(mh, s[lane * QCH + c] * 0.0625f); + for (int c = 0; c < QCH; ++c) { + const float e = c < nc ? exp(s[lane * QCH + c] * 0.0625f - mh) : 0.0f; + lh += e; + p[lane * QCH + c] = half(e); + } + } else if (lane < 16) { + for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); + } + simdgroup_barrier(mem_flags::mem_threadgroup); + if (qsa_stop == 3) { if (lane < HPG) { pp[lane * (QD + 2) + QD] = 3.0f; pp[lane * (QD + 2) + QD + 1] = 1.0f; } return; } + // ---- acc = P V, 64 dims at a time + { + // K = 32 cells a run (two runs a 64-cell chunk): K = 64 runs compute garbage on the M4 + constexpr auto desc = matmul2d_descriptor(16, 64, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tP(p, dextents(QCH, 16)); + tensor, tensor_inline> tO(pp, dextents(QD + 2, HPG)); + for (int d0 = 0; d0 < QD; d0 += 64) { + auto rv = op.template get_right_input_cooperative_tensor(); + auto tPh = tP.template slice<32, 16>(0, 0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int c0 = 0; c0 < QCH; c0 += 32) { +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + const uint row = uint(max(id, 0)) * QKV + g; // clamped, as in the S fill above + const float v = float(vq[row * QD + d]) * float(vs[c][d >> 6]); + rv[i] = id >= 0 ? half(v) : half(0.0f); + } + tPh = tP.template slice<32, 16>(c0, 0); + op.run(tPh, rv, acc); + } + if (qsa_stop == 4) { // the bisection: the old direct store into the 258-float rows + auto mO = tO.slice(d0, 0); + acc.store(mO); + continue; + } + // through threadgroup memory (s is free now): a matmul2d store into rows of QD + 2 floats loses the + // whole threadgroup on the M4 + tensor, tensor_inline> tT(s, dextents(64, 16)); + simdgroup_barrier(mem_flags::mem_threadgroup); + acc.store(tT); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < HPG * 64; it += 32) pp[(it >> 6) * (QD + 2) + d0 + (it & 63)] = s[it]; + } + } + if (lane < HPG) { pp[lane * (QD + 2) + QD] = mh; pp[lane * (QD + 2) + QD + 1] = lh; } +} + +// qcur (fp32) -> fp16 for qsa_attn_mm +kernel void qsa_q_f16(device const float* q [[buffer(0)]], device half* qh [[buffer(1)]], constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (gid < n) qh[gid] = half(q[gid]); +} + +// qsa_score as a GEMM (fp32 operands, so the selection sees qsa_score's values to fp32 rounding): S[(t, h)][b] = +// qidx[t][h] . pooled[b] for 16 tokens x IH heads (64 rows) by 64 blocks a threadgroup; the epilogue sums relu over +// the heads. The tail block n_bid (the `dead` key, + 1e9 when it has cells) and blocks past it are left to the +// scalar path: the tail is scored here by one simdgroup a token; blocks > n_bid are not written (qsa_topk reads +// b <= n_bid). grid (ceil(max blocks / 64), ceil(T / 16)), 128 threads. +kernel void qsa_score_mm(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]], + device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]], + constant QsaArgs& a [[buffer(4)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float ot[64 * 64]; + const int b0 = int(tg.x) * 64, t0 = int(tg.y) * 16; + const uint nb_all = (a.pos0 + a.n_tok) / 4 + 1; // blocks any token of the chunk can score + tensor, tensor_inline> tQ((device float*)qidx, dextents(ID, int(a.n_tok) * IH)); + tensor, tensor_inline> tK((device float*)pooled, dextents(ID, int(nb_all))); + tensor, tensor_inline> tO(ot, dextents(64, 64)); + constexpr auto desc = matmul2d_descriptor(64, 64, dynamic_length_v, false, true, false); + matmul2d> op; + auto mA = tQ.slice(0, t0 * IH); + auto mB = tK.slice(0, b0); + op.run(mA, mB, tO); + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint stride = a.max_cells / 4 + 2; + for (int idx = tid; idx < 16 * 64; idx += 128) { + const int tt = idx >> 6, j = idx & 63, t = t0 + tt, b = b0 + j; + if (t >= int(a.n_tok)) continue; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + if (uint(b) >= n_bid) continue; + float s = 0.0f; + for (int h = 0; h < IH; ++h) s += max(ot[(tt * IH + h) * 64 + j], 0.0f); + scores[uint(t) * stride + b] = s; + } + // the tail block of each token in this threadgroup's block range (the first column threadgroup does it) + if (tg.x == 0) + for (int tt = sg; tt < 16; tt += 4) { + const int t = t0 + tt; + if (t >= int(a.n_tok)) break; + const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2; + float k[4]; + for (int i = 0; i < 4; ++i) k[i] = dead[lane + 32 * i]; + float s = 0.0f; + for (int h = 0; h < IH; ++h) { + float d = 0.0f; + for (int i = 0; i < 4; ++i) d = fma(qidx[(t * IH + h) * ID + lane + 32 * i], k[i], d); + d = simd_sum(d); + s += max(d, 0.0f); + } + if ((n_kv & 3) != 0) s += 1e9f; + if (lane == 0) scores[uint(t) * stride + n_bid] = s; + } +} + +// qsa_flash: the whole attention of a (kv head g, token t) in one threadgroup (no split-K partials, no merge): the +// width selected cells in chunks of 64 with an online softmax; 4 simdgroups, simdgroup s owns dims 64 s .. 64 s + 63 +// (its slice of S's K sum and its 64 output dims). Out: xq = (acc / l) * sigmoid(gate) as fp16 (attn_output's GEMM +// operand). grid (QKV, n_tok), 128 threads. +kernel void qsa_flash(device const half* qh [[buffer(0)]], device const char* kq [[buffer(1)]], + device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]], + device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]], + device const float* qfull [[buffer(6)]], device half* xqh [[buffer(7)]], constant QsaArgs& a [[buffer(8)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int HPG = QH / QKV; + threadgroup float sp[4][16 * QCH]; + threadgroup half p[16 * QCH]; + threadgroup int cid[QCH]; + threadgroup float vs[QCH]; + threadgroup float fac[16], lsum[16], mrun[16]; + const int g = int(tg.x), t = int(tg.y); + const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET)); + const int nch = int((width + QCH - 1) / QCH); + if (tid < 16) { mrun[tid] = -INFINITY; lsum[tid] = 0.0f; } + constexpr auto dS = matmul2d_descriptor(16, QCH, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + constexpr auto dO = matmul2d_descriptor(16, 64, 32, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d opS; + matmul2d opO; + tensor, tensor_inline> tQ((device half*)qh + (uint(t) * QH + g * HPG) * QD, + dextents(QD, HPG)); + tensor, tensor_inline> tP(p, dextents(QCH, 16)); + auto rv = opO.template get_right_input_cooperative_tensor(); + auto tPh = tP.template slice<32, 16>(0, 0); + auto acc = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + // (a second destination tensor, not a copy of acc: the MetalPerformancePrimitives headers of macOS 26.3 cannot copy a + // cooperative tensor - "cannot initialize a parameter of type 'void *' with an rvalue of type 'const void *'", #2323; + // tot is zeroed right below, so the copy was only there for its type) + auto tot = opO.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); // the running output; acc takes one chunk (fresh each chunk: on the M4 a run does not accumulate + // into elements rescaled in place) + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = 0.0f; + const int d0 = sg * 64; + for (int ch = 0; ch < nch; ++ch) { + const int c0 = ch * QCH, nc = min(QCH, int(width) - c0); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < QCH) { + const int id = tid < nc ? ids[t * QBUDGET + c0 + tid] : -1; + cid[tid] = id; + vs[tid] = id >= 0 ? float(vsc[(uint(id) * QKV + g) * 4 + sg * 0]) : 0.0f; // (per-simdgroup scale below) + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // S partial over this simdgroup's 64 dims + { + auto rk = opS.template get_right_input_cooperative_tensor(); + auto tA = tQ.template slice<32, 16>(d0, 0); + auto s = opS.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, + __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < s.get_capacity(); ++i) + if (s.is_valid_element(i)) s[i] = 0.0f; + for (int k0 = d0; k0 < d0 + 64; k0 += 32) { + if (!coop_k4) { + for (uint16_t i = 0; i < rk.get_capacity(); ++i) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], id = cid[ix[1]]; + const uint row = uint(max(id, 0)) * QKV + g; // clamped, as in qsa_attn_mm + const float v = float(kq[row * QD + k]) * float(ksc[row * 4 + (k >> 6)]); + rk[i] = id >= 0 ? half(v) : half(0.0f); + } + } else +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rk.get_capacity(); i += 4) { + const auto ix = rk.get_multidimensional_index(i); + const int k = k0 + ix[0], c = ix[1], id = cid[c]; + half4 w = half4(0.0h); + if (id >= 0) { + const uint row = uint(id) * QKV + g; + const char4 q4 = *(device const char4*)(kq + row * QD + k); + w = half4(float4(q4) * float(ksc[row * 4 + (k >> 6)])); + } + rk[i] = w.x; rk[i + 1] = w.y; rk[i + 2] = w.z; rk[i + 3] = w.w; + } + tA = tQ.template slice<32, 16>(k0, 0); + opS.run(tA, rk, s); + } + tensor, tensor_inline> tS(sp[sg], dextents(QCH, 16)); + s.store(tS); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // online softmax per head row: simdgroup 0, lane = row + if (sg == 0 && lane < 16) { + if (lane < HPG) { + float m = mrun[lane]; + for (int c = 0; c < nc; ++c) { + const int q = lane * QCH + c; + m = max(m, (((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f)); + } + const float f = exp(mrun[lane] - m); + float l = lsum[lane] * f; + for (int c = 0; c < QCH; ++c) { + const int q = lane * QCH + c; + const float e = c < nc ? exp(((sp[0][q] + sp[1][q]) + (sp[2][q] + sp[3][q])) * 0.0625f - m) : 0.0f; + l += e; + p[q] = half(e); + } + mrun[lane] = m; + lsum[lane] = l; + fac[lane] = f; + } else { + for (int c = 0; c < QCH; ++c) p[lane * QCH + c] = half(0.0f); + fac[lane] = 0.0f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // rescale this simdgroup's accumulator rows, then acc += P V over its 64 dims + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int c0 = 0; c0 < QCH; c0 += 32) { // K = 32 cells a run (K = 64 is wrong on the M4) +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < rv.get_capacity(); ++i) { + const auto ix = rv.get_multidimensional_index(i); // (cell, dim) + const int c = c0 + ix[0], d = d0 + ix[1], id = cid[c]; + const uint row = uint(max(id, 0)) * QKV + g; // clamped, as in qsa_attn_mm + const float v = float(vq[row * QD + d]) * float(vsc[row * 4 + sg]); + rv[i] = id >= 0 ? half(v) : half(0.0f); + } + tPh = tP.template slice<32, 16>(c0, 0); + opO.run(tPh, rv, acc); + } + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) tot[i] = fma(tot[i], fac[acc.get_multidimensional_index(i)[1]], acc[i]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) { + const auto ix = acc.get_multidimensional_index(i); // (dim, row) + const int h = ix[1], d = d0 + ix[0]; + if (h < HPG) { + const int hh = g * HPG + h; + const float gate = qfull[(uint(t) * QH + hh) * 2 * QD + QD + d]; + xqh[(uint(t) * QH + hh) * QD + d] = half(tot[i] / lsum[h] * (1.0f / (1.0f + exp(-gate)))); + } + } +} diff --git a/edge0-strata/metal/kernels/quant.metal b/edge0-strata/metal/kernels/quant.metal new file mode 100644 index 00000000..c665f503 --- /dev/null +++ b/edge0-strata/metal/kernels/quant.metal @@ -0,0 +1,701 @@ +// metal/kernels/quant.metal - the weight formats as "units": a unit is UNIT consecutive weights of one row that one +// SIMD lane decodes at a time into raw values q[0..UNIT) and one scale s, so that the unit's dot with x is +// s * sum(q[j] * x[j]). Every GEMV kernel (moe.metal, gemv.metal) is written once over these. +// +// Decode rules transcribed from ggml-quants.c dequantize_row_* (3cf03257) and Strata's Q2_0 contract +// (docs/q2_0-contract.md); the grids and sign tables are ggml-common.h's (prepended by the runtime). The tables a +// format needs are copied to threadgroup memory once per threadgroup (constant-memory lookups with divergent indices +// serialise on Apple GPUs). +// +// Blocks start at 2-byte aligned addresses (18, 66, 74, 82, 98, 110 bytes), so every multi-byte field is read as +// ushorts. + +#define T_Q4_0 2 +#define T_Q3_K 11 +#define T_BF16 30 +#define T_F32 0 +#define T_Q5_0 6 +#define T_Q5_1 7 +#define T_Q8_0 8 +#define T_Q4_K 12 +#define T_Q5_K 13 +#define T_Q6_K 14 +#define T_IQ2_XXS 16 +#define T_IQ2_XS 17 +#define T_IQ3_XXS 18 +#define T_IQ4_NL 20 +#define T_IQ3_S 21 +#define T_IQ1_S 19 // Unsloth UD-IQ1_S / UD-IQ1_M experts (50-byte blocks) +#define T_IQ1_M 29 // the 2.00 bpw pack's 1.75-bit type +#define T_IQ2_S 22 +#define T_IQ4_XS 23 +#define T_Q2_0 42 +#define T_Q2_0R 142 // Q2_0 read with raw activations (codes - 1 in the decoder; no prepared x needed) +#define T_TQ2T 144 // agentionai's trellis types (Gyro): 128 values a block, 34 / 26 / 30 bytes, Hadamard-folded experts +#define T_TQK6 145 +#define T_TQK7 146 + +template struct fmt; +// QK block elements, BB block bytes, UNIT weights per unit, TAB words of threadgroup table +template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 64, BB = 18, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 98, UNIT = 16, TAB = 256 + 32 }; }; +template <> struct fmt { enum { QK = 256, BB = 66, UNIT = 16, TAB = 512 + 32 }; }; +template <> struct fmt { enum { QK = 256, BB = 56, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 50, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 74, UNIT = 16, TAB = 1024 + 32 }; }; +template <> struct fmt { enum { QK = 256, BB = 82, UNIT = 16, TAB = 2048 + 32 }; }; +template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 16, TAB = 512 + 32 }; }; +template <> struct fmt { enum { QK = 32, BB = 18, UNIT = 32, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 136, UNIT = 32, TAB = 256 }; }; +template <> struct fmt { enum { QK = 256, BB = 110, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 144, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 176, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 256, BB = 210, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 16, BB = 32, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 16, BB = 64, UNIT = 16, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 34, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 18, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 22, UNIT = 32, TAB = 1 }; }; +template <> struct fmt { enum { QK = 32, BB = 24, UNIT = 32, TAB = 1 }; }; +// trellis: a unit is one lane of 4 steps = 16 weights; TAB = the 2048-point codebook as half2 words +template <> struct fmt { enum { QK = 128, BB = 34, UNIT = 16, TAB = 2048 }; }; +template <> struct fmt { enum { QK = 128, BB = 26, UNIT = 16, TAB = 2048 }; }; +template <> struct fmt { enum { QK = 128, BB = 30, UNIT = 16, TAB = 2048 }; }; + +static inline uint rd32(device const uchar* p) { + device const ushort* s = (device const ushort*)p; + return uint(s[0]) | (uint(s[1]) << 16); +} +static inline float rdh(device const uchar* p) { return float(*(device const half*)p); } + +// ---- threadgroup tables ------------------------------------------------------------------------------------------- +// layout in `tab` (uint words): [0, G) the grid, [G, G + 32) ksigns_iq2xs as bytes; IQ4_NL/XS: 16 floats. +template static inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth); + +// QUANT_CTAB (function constant 23): the i-quant grids and the sign table read from ggml-common's constant arrays +// instead of a threadgroup copy (no table load and no barrier wait at the threadgroup's start); the same values. +constant bool QUANT_CTAB [[function_constant(23)]]; +constant bool ctab = is_function_constant_defined(QUANT_CTAB) && QUANT_CTAB; +static inline uint ksg(threadgroup const uint* tab, int at, uint i) { + return ctab ? uint(ksigns_iq2xs[i]) : uint(((threadgroup const uchar*)(tab + at))[i]); +} +static inline uint g_iq3xxs(threadgroup const uint* tab, uint i) { return ctab ? iq3xxs_grid[i] : tab[i]; } +static inline uint g_iq3s(threadgroup const uint* tab, uint i) { return ctab ? iq3s_grid[i] : tab[i]; } +static inline uint g_iq2xxs(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2xxs_grid[i] >> (32 * w)) : tab[2 * i + w]; } +static inline uint g_iq2xs(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2xs_grid[i] >> (32 * w)) : tab[2 * i + w]; } +static inline uint g_iq2s(threadgroup const uint* tab, uint i, int w) { return ctab ? uint(iq2s_grid[i] >> (32 * w)) : tab[2 * i + w]; } + +// BEGIN tq_lut (generated by metal/tools/gen_tq_lut.py from third_party/ggml/ggml-common-tq.h; do not edit) +constant uint tq_lut_u[2048] = { + 0xb8f7bd4d, 0x3cd5b6ce, 0x4368bd9e, 0xb8a63dce, 0xb430395e, 0x3dbc3bc1, 0x3c9cb92b, 0x3aaabca2, + 0xb8e13cdb, 0x3dfbbb33, 0xbb183a8d, 0x3da0bd64, 0x3b83bf1a, 0x3659b389, 0xabd7b40f, 0x35a240ee, + 0x3972c057, 0x3996be17, 0xb9b5b926, 0x3c75393b, 0xbdef4151, 0xbd09b5a3, 0x4012b408, 0x3ece3b7f, + 0x4011b688, 0xbe303b17, 0xb4bf3d9b, 0xb8c8b870, 0x38b2b537, 0x34b031be, 0x38713389, 0x38b8ba33, + 0x3b3dbce0, 0x3d3d3367, 0x3cbe3f50, 0xb99130b0, 0x39863733, 0xb6d9b70e, 0xbd98be8c, 0xbe04bcca, + 0x403c36a9, 0x3eaeb527, 0x39823c6a, 0xbd5bb655, 0x35963ebb, 0xbbbd417e, 0xb1d32fe0, 0x31c0b858, + 0xbc10b54e, 0xaf8b3cf0, 0x3f5f38e6, 0xb2f63b87, 0xb675b114, 0x329829bf, 0xb955bb7b, 0xa98434e9, + 0x378fc04a, 0x3c8c31a0, 0x3c682dea, 0x2d292fe1, 0xbd722c33, 0x3e3e379d, 0x3cf2b918, 0x38bfbcbe, + 0xbdd53284, 0xbd013707, 0x38c1bec1, 0x3838b960, 0x30103413, 0x3672376d, 0x3f42302f, 0xbe9a3a61, + 0xbc713c06, 0xbcf2bb57, 0x3c85a9dc, 0xb8b23873, 0x340d2d1e, 0xb73f3661, 0xb8b7bb91, 0x3943b956, + 0x39853958, 0xb5563b26, 0xc2b9ae17, 0xb9ea34e9, 0xb6f433d5, 0xb1af3826, 0xaf98b8a4, 0xbc693a2d, + 0x358e3a25, 0x32fd406d, 0xb11a38a4, 0xb02eb170, 0xb5cbaa0a, 0x3d453c04, 0xb55bb258, 0xbbf5b412, + 0x3f4b3d33, 0x38a9bc6e, 0xb715b8c0, 0xbdd92ec6, 0x2980beb3, 0xbb3c3c54, 0x2d683ba0, 0xb20942f9, + 0x3337b557, 0xbda73870, 0x35bd2d57, 0x385eb8cd, 0xb68e39fe, 0xbd743eda, 0xbd543643, 0x37d43a32, + 0x2f7f38f8, 0x380f37f3, 0xbc231b1e, 0x35c825c7, 0xb82a2e3e, 0x3897b85f, 0x3cb33c0c, 0x377bc13e, + 0x3f09b3a4, 0x3c73c006, 0x3ad1bda5, 0xa3cfb0c6, 0xb680b9d5, 0xb493384b, 0xa9853a3c, 0xba24bcb5, + 0x3c21bd0c, 0x34dc2f68, 0xbdb7bc5d, 0xbddbbbbf, 0xb95ab4c4, 0xbdaab59b, 0xbdce34cf, 0x3f9d3bad, + 0x3bd53a05, 0x2c64b69c, 0x3adb4003, 0x35e83d4c, 0xb5cc308d, 0xbb86bd2d, 0x3cd3b124, 0xb2ab3a0a, + 0x3d4bb507, 0x284cb190, 0xbe83384c, 0xb806bc94, 0xb14631e4, 0x3178bbe0, 0x3819bc9d, 0x2e8f3824, + 0x4150c246, 0x3a43bcf3, 0xbdd53da9, 0x3c47b229, 0xc001400e, 0xb0c33bb2, 0x3822bd46, 0xa4a9b92d, + 0x3cfa3c72, 0x42843550, 0xbc97b9e5, 0x357ab996, 0xaf70ba03, 0xbc0735e7, 0x397bb663, 0x29423c59, + 0xb689adf7, 0x3db8b875, 0xb63c35ea, 0x33493674, 0x3ae03926, 0xb56ebb20, 0xa9ae3f05, 0x3c383ce5, + 0xb1ef40b8, 0xba2bba73, 0x3d3fbdd7, 0x3c093822, 0xbce4bd87, 0x36beba38, 0xbe23b808, 0xbd77ba33, + 0x34663b90, 0xb805c1c5, 0x375b34c0, 0xa2dbb4be, 0x3495b96c, 0xb99ab38a, 0x31d5390a, 0x3e16a910, + 0x3dd62878, 0xbdc03964, 0xb3efb783, 0xb571ba12, 0xc0f2b8ce, 0x3771bd8f, 0x3eba3656, 0x3c39be0e, + 0x371b3ac5, 0x40713469, 0x2e19bd99, 0x3cb83ab2, 0xa92c36d9, 0xbc9bbcb4, 0xb4e2a0fa, 0x31fd34ed, + 0xb0a04033, 0x342ebac8, 0xbccd3df5, 0xbc5fbc5a, 0x3ada3f32, 0x39ec3878, 0x3ce53db1, 0xbe93b891, + 0xb75e3c11, 0xbb27b82c, 0xba6736d6, 0x352cb602, 0xbf3a3507, 0xbcf5b8e9, 0x3773b55e, 0xbd0ab832, + 0x3989bf01, 0x3d733edf, 0x3bc8bb3d, 0xb7e0bf82, 0x393f4053, 0x2f91b07d, 0xb6afbc8b, 0x2a463538, + 0xbbfd3c2f, 0x324ebfc9, 0xb7023a8f, 0xbda9b939, 0x3c3f3ac3, 0x37983875, 0x3e1b3dc4, 0x255eb820, + 0xbe92a7e2, 0xb7c7b63f, 0xbcb438bd, 0xb9583a2f, 0x34fcbdf2, 0xb5a83ba1, 0x30473975, 0x21daaecc, + 0x3bc4b7f7, 0x2eb639f8, 0xbb3abb3f, 0xb13dbd9b, 0xbab237fa, 0x30743794, 0xbad3bd7e, 0x346b3d65, + 0xba613cfd, 0xb8feadd5, 0xb6ecbe09, 0xb544c05b, 0x2fdb360c, 0x39e5b8ba, 0x3e96b8a8, 0x38cf40cf, + 0xac5aba8b, 0xb284b691, 0xb5d5c0b9, 0x3d003864, 0xa3823220, 0xbb8f3422, 0xb85eba85, 0x3eb33854, + 0xb885c110, 0x2e60bdf8, 0xb8ba1cb3, 0x3636b4d3, 0xaf93382a, 0x346abc81, 0x31483d66, 0x36fec0b3, + 0x4001af70, 0x3eac3caf, 0xb614bd48, 0x3de83d04, 0x31d03f2b, 0x30ceb955, 0x2d33b81b, 0xb657b328, + 0x3527b7c2, 0xc255b9f6, 0x3c90bc1a, 0x35f3be9c, 0x409ac11d, 0xbf2f30db, 0x3285be02, 0x2f14b607, + 0x35c23c56, 0x30ac3288, 0xc25c3c1c, 0xbbd63ca4, 0x29c83465, 0x412937a3, 0x32253151, 0xad412ef2, + 0xbfc1bbcb, 0xbe71b0c0, 0xbb59b17d, 0xbefb3cfe, 0x39f0bdb8, 0x28a74014, 0x396ab771, 0xc063bbdc, + 0x2898babe, 0x3de9b683, 0xaa78bd1f, 0x3ded4187, 0x3d7b3e18, 0xab0caf3c, 0xb52438a7, 0xc38fb964, + 0x3d02beda, 0xb25daa74, 0xbe9a2d90, 0xbbb0be23, 0x3a8cbf6a, 0x3d87bbfa, 0x3571b53c, 0xc027385c, + 0xb801b903, 0x3b7b3368, 0x3ab8b70c, 0xb62538b5, 0x3cb1b436, 0x3f2c8d2a, 0x2c53b88e, 0x393bb2c8, + 0x3f40bdbf, 0x39d9bc1b, 0x306dbc3a, 0xb8fcb3ba, 0x3c75c07b, 0xb421b655, 0xb7203d34, 0xb6a037f1, + 0xbec7b338, 0x334eba58, 0xb8103aab, 0xbe3db3d7, 0x40a2b984, 0xba54b5cd, 0xb63e3afa, 0x394dbd80, + 0x2ddaba54, 0xb8e4baec, 0xad382b2a, 0x3886bd82, 0xc015ac6b, 0x3c72368b, 0xbaa72fb3, 0xb2cb38fb, + 0xbfb5335d, 0xa535aba9, 0x38b7afb8, 0x35243d00, 0xaeac3b42, 0xb1a64144, 0xba3ab352, 0x39f32c72, + 0xb82c4079, 0xb45fb9e5, 0xb9f59d20, 0xbd36b2ee, 0x2f11bad5, 0x3e803450, 0xb684213f, 0xbd073200, + 0xbc5eb839, 0x3f5dc217, 0xbe943d9a, 0x33813947, 0xb9439081, 0x37953eac, 0x3915ad61, 0x3d5db810, + 0x33f239cb, 0x41edbfe8, 0xb4a23ee7, 0xb88ab9e5, 0xc0da36b3, 0x406cb57f, 0x3c5cba70, 0x3ae4b936, + 0xb7ff3e2c, 0xb08e35b6, 0x3c3ab8d1, 0xafed227e, 0x310a3a3d, 0xb9262d00, 0x3f76bbe1, 0x3e15b8aa, + 0xc074b31c, 0xbcb640ee, 0x3b033449, 0xbd0138f5, 0x386abc2d, 0xbcfc3f94, 0xb04933a2, 0xb86e3c77, + 0xb91c366c, 0x35c6b832, 0xbc4c36ab, 0x3a04b6b9, 0x3c84b2c9, 0xb7623006, 0x3580330e, 0x39263d32, + 0xbfb238f9, 0x3de6c00a, 0xafad3a35, 0x3d88be69, 0xbe12c30f, 0x3d26b13d, 0xb7a5382a, 0x3ac340fa, + 0xc028b8a9, 0x3e19b979, 0xbdd7bac1, 0xb802bd42, 0xbe1ab631, 0xac7ab8fb, 0x38193ff8, 0x42e83d3f, + 0xb1dcb312, 0xaa42c036, 0xc078bd28, 0x36ddb6f5, 0x31fdbad4, 0xb7e93933, 0x31d439b4, 0xb4aebd13, + 0x3c5d37d3, 0xbafd3925, 0x38753b52, 0xa5812a38, 0xb4503b51, 0x3cccc248, 0x3ad2b1aa, 0xbe9036eb, + 0x34e3423c, 0x3bc13ac2, 0x37493be5, 0x3ceebcd1, 0x40173ac7, 0xbfc43a67, 0xb9a8b7eb, 0xbf08bf94, + 0xa983b711, 0xb6093c0f, 0xb7f73ebd, 0xbca734f4, 0x38b1bac3, 0xb93fb5d0, 0x3f1d3e04, 0x305cb7fd, + 0xbcaaaae3, 0x3d8cbb0f, 0xbd3fbca4, 0xbc08bb3b, 0xb9de3a90, 0xad25377f, 0xbefeb988, 0xbc6cbdf7, + 0x35c3ba23, 0x370abb14, 0xb3a93c45, 0x3c72b889, 0x3e713a36, 0xb8e03a90, 0x34a634cb, 0x3c6c40be, + 0xbc3fbe95, 0x2284ba4d, 0xb73eb96e, 0xbc04384d, 0x3a03b452, 0xc10041a8, 0x2ed23c75, 0x3ab6bc47, + 0x3a633746, 0x3735b62f, 0x3c2fb9dd, 0xb08fbba8, 0x3289b270, 0x3c1e3972, 0xb5ae3844, 0x3d303cee, + 0x3fa2b98f, 0xb1503e16, 0x382c3c9c, 0xae2e3483, 0x3a843b71, 0x32f4bb53, 0x3ab23540, 0xba06b6e0, + 0x3701bcc9, 0x361ab89a, 0x3a74be12, 0x37ae3b56, 0x38eeb6ec, 0xb60ab47a, 0x374a3247, 0x3883394d, + 0x2c41b76f, 0xb862bc07, 0xbba5395e, 0x3865b71f, 0x2cbc3897, 0x35c9ac7a, 0x3ac5b9df, 0xb8b9b63b, + 0x3ff6c064, 0xb8293ba3, 0x3b7eb904, 0xb78eafdb, 0x3c22baec, 0x39a5b53d, 0xc0d4bab6, 0xbddbb47e, + 0xc01cba66, 0xbcb037ec, 0x3ac3b409, 0x3e472b2b, 0xb8e03b8a, 0x3a7738cd, 0xb7243cb0, 0x381fbaef, + 0xb0f9ba7a, 0xb1383940, 0xb8403855, 0x35f439af, 0x2c193b0f, 0x3b833760, 0x3801a813, 0x41d6b801, + 0x3ab33026, 0xa9f5bca3, 0x370ebf1e, 0x3467b82b, 0xbd57b4c3, 0xb4e9ad4c, 0xbad6b8c4, 0xbd66a8f7, + 0x3cd23e70, 0x3421236c, 0x3c22aa8a, 0x9661b9c2, 0xb6223988, 0xb8aab797, 0x305c254c, 0x271bbd53, + 0xbaa93c11, 0x3601b75b, 0xb4e936d1, 0xb15db856, 0xc0b2be96, 0x3e32b7c5, 0x3e22b069, 0x3b41bc18, + 0x31cac16e, 0x35f9365a, 0xb700b491, 0xb41b3a6a, 0x406a2f26, 0xbd54b8a6, 0x39e33b9b, 0xad0fbeae, + 0xb1663b1b, 0x3655afd4, 0xbaa7bde7, 0xbd6c37c1, 0x3cbfbc6d, 0x2dbd3584, 0xb9f93d69, 0xac53b992, + 0xb2a2bf6b, 0xc15735b5, 0x328e331a, 0xb7b3ba14, 0x38093c44, 0x33613c08, 0xbc1c3b93, 0xbd213b53, + 0x39b2b9d1, 0xbad9326b, 0x35e6b1c8, 0xaeff38ff, 0xb734b831, 0xbbe1ba96, 0x36fb33f8, 0xc15ab049, + 0xb7a33b22, 0xbc3db462, 0xb57cbf97, 0xbcaa3c86, 0x3c6ebcb9, 0xb0f539ca, 0xb5bf379f, 0xbcb6b629, + 0x3f0f3aab, 0x34fcc01a, 0xb6be3869, 0x32f4b953, 0xbdb7b85f, 0x411a39f4, 0x203cbce8, 0x2a593824, + 0xa60cb26f, 0x3d96b47f, 0xb53fb8d0, 0xb69f3706, 0x36cfb1c9, 0x37c1ad03, 0xbc903d80, 0xb98eba26, + 0x334b3c5b, 0xb051c196, 0xbd79bd24, 0x3e2f3b38, 0xb64dba82, 0x34ed2a56, 0x2a01b36f, 0xbfbc40e9, + 0xb600b7ce, 0x3a9c3cce, 0xb23c33c6, 0x36c83c33, 0x409ea776, 0x3c35b822, 0x30ef3e43, 0xb9f7b9b2, + 0xba813c79, 0xaedb3973, 0xba33bfe1, 0x3d8a348e, 0xbfa1b5ac, 0x31b7b8d9, 0x34813aee, 0xb56e3d3b, + 0x2c653d98, 0x3d7439e0, 0x3ae2b51a, 0x3bb9ae6c, 0x1d86bdc8, 0x3840b9d9, 0xc0f1bc60, 0x396d3e26, + 0xa4783dc5, 0xc01bbf35, 0x39c8b938, 0x3fd035d1, 0x3d41b3bc, 0x3bf9bca8, 0x32f5bd36, 0x34c83f59, + 0x38543e38, 0xb82e338e, 0xa3ce37b2, 0x37fdb683, 0x3f02c104, 0x38b3b92f, 0x3bcdb454, 0xba4231f4, + 0x3878bf4d, 0x3559bf7f, 0xbc53b8e0, 0x407ebd77, 0xb0c6b69e, 0xbabdacce, 0xb92f3c1b, 0x3a5fb8fa, + 0x3f13b628, 0xbacd3893, 0x3ec8c047, 0xba7fb70b, 0xb140b997, 0xb4662c0c, 0x3ddeba3f, 0x3ae1be75, + 0x3879b466, 0x3c553d68, 0x3c6eb4d9, 0xb803be71, 0x2a523e15, 0xbc373e1b, 0xb7473585, 0x3db4bf22, + 0x31a6bcdb, 0x2e23b415, 0xb884ab25, 0xb5ef329e, 0xb59c3525, 0x39c2ba76, 0xbf1eb6f5, 0x2dadbcdf, + 0x40193e8c, 0xbb4ab92c, 0x2ddb3494, 0x3afd406e, 0xb97a3b91, 0xbce13c1f, 0x3c343eae, 0x38ecbddd, + 0x3e4f390f, 0x3e04bc20, 0xba0b378c, 0xb8aa330d, 0x348e3e34, 0x352fbaa5, 0x3d2bbac3, 0xbe71b982, + 0xc1e1305c, 0xbe12be05, 0x2fc1b70b, 0xbe713c8b, 0x3dcc35ad, 0xb612b581, 0xb50cb42c, 0xb36cb1c9, + 0x3977b443, 0x3d11bd4c, 0x3b1739b8, 0xbd27269a, 0x3d6cbce0, 0xb611b86d, 0xb50c35e7, 0x359cb453, + 0x38bf38c6, 0x39a6b82a, 0x3bb1a8c2, 0x3c8bb020, 0x2dfa2c82, 0xb43ebb37, 0xbbbd2b23, 0xbd253c97, + 0xbb813d21, 0xbe12b8e5, 0x3d6a3845, 0xc0ecb571, 0x34ba3c69, 0x2818b96f, 0x382dbdde, 0x3f333f0c, + 0x3c30b44b, 0xbef5aebd, 0x26833a3b, 0xb7f8aaf0, 0xb4f5405e, 0xb1973c5b, 0xbada359f, 0xbb773227, + 0x3dec3877, 0xb7c634a2, 0xba0e3bde, 0x20ccb758, 0xb8e83760, 0x3cb8bdaf, 0x385d3cf3, 0x3078306f, + 0x34d1ac7a, 0x402fb9b8, 0xbc802d2e, 0xbeefb55d, 0x356a348e, 0x3b683511, 0x32ea35a4, 0x3a133e99, + 0xb8c94002, 0x355a3dc1, 0xbb783556, 0x37ee3560, 0xbc4aba6e, 0x33ebb14d, 0x3c4a22a9, 0x3a27bd5b, + 0xba1dc0ca, 0xbd314053, 0xb87b3c1a, 0x36e338de, 0x2b952663, 0x36853b4f, 0xc020353c, 0x4050be97, + 0xa53abc48, 0xbeea3e4e, 0xbb52ac24, 0x40d2b743, 0x3762b389, 0xbb02b493, 0x342f3cb4, 0x3aceba83, + 0x3923bc28, 0xc0afc004, 0xb842b898, 0x276038d0, 0x405eb83e, 0x4254affb, 0x3c6f3b69, 0x3bedba70, + 0xba70b0de, 0xb6353d9c, 0x3de8bde3, 0x38e53321, 0x40ae369f, 0xbaf236e3, 0xbb9baf7b, 0xc1f3b484, + 0x40973f5a, 0xaead3da5, 0x1e3a3934, 0x409fbc67, 0x3a55b2ed, 0x385a3009, 0xb6ac3166, 0xc3f83832, + 0xa95fb800, 0x2f61bc8a, 0xbeba353c, 0x3ce33284, 0xade83fad, 0xb7c13da0, 0x382db56e, 0x3855385f, + 0x2ff5bd34, 0xb290ba07, 0xb6f4bda4, 0x3ad2b89b, 0xbe273613, 0x35fc3838, 0x3c9628e1, 0xbd623c1e, + 0x356b315a, 0x3b163b05, 0x3ed5baf3, 0x3c65bf42, 0x378a3d3f, 0xbd43bf4e, 0x39643c0d, 0x3c0034c3, + 0xbafe345a, 0xba8bbf4f, 0xb95b4240, 0xc0a23b21, 0x331db6ed, 0x3111b57a, 0xb8ca3594, 0x3d96c16a, + 0x3c86b7a9, 0xb498b949, 0x3c1e2d89, 0x2eebad06, 0xba8abc5b, 0xb0b9be54, 0x3e71bc9a, 0x3ff43ca7, + 0xbbbcbf8b, 0x3fefb894, 0x40c93bda, 0xbb663bc7, 0x3cbbb87e, 0xb1f7bb1e, 0x38f63a3d, 0x3cd9ad22, + 0xbd81bdb4, 0xbd96b005, 0xb91c3d6e, 0xbd5e38f2, 0x2e27b513, 0xbe1c3e5b, 0x366e34b3, 0xb6cdbb24, + 0xbfb436f7, 0x34a3b00c, 0x37dfb1bf, 0x33d3b9d6, 0x3edf396f, 0xb6f7b5d1, 0x38b22aa1, 0x223d36b1, + 0x3d3ebc6a, 0xb7b5b777, 0xb91cb8ed, 0x35f53be6, 0xbd74c077, 0x3966bcdc, 0x3d0bbc12, 0x28cb3312, + 0x4128bd0d, 0xb80ebef1, 0x3eb4b758, 0x352fb6d3, 0x3ca3343d, 0xbaefb733, 0xbc1dbce4, 0x38fa3eac, + 0x3bceb9d1, 0x36ec3953, 0x3985bbb4, 0x3fe24005, 0xb9913f4e, 0xbe2c39d5, 0xb37fb44e, 0x30aa3acf, + 0x40c3b358, 0x3b9b3c1f, 0x39a739ec, 0x41113d01, 0xbb76367d, 0x3a483c75, 0x3c01b547, 0x3ae03d2f, + 0xb837c3b1, 0x36c12e1d, 0xbd8eb2d3, 0x3d80b23a, 0x3c9bbd29, 0x348c37a8, 0xb5aa3c9e, 0x40d33311, + 0x4028256d, 0xa66e2e47, 0xb8de390a, 0x3c033306, 0xbb3fbc58, 0xbc21bc05, 0x28e2376c, 0x3ddab536, + 0x3ebb2427, 0xbdbfb70a, 0xbb103cc4, 0x2d293d33, 0xbfd0c044, 0x3726bfd3, 0xad3432d9, 0x3d32c040, + 0x2b123f66, 0x3ae9377a, 0xc02d3088, 0x38e2384f, 0xacfb3c2b, 0xb9debbda, 0x39452d0a, 0x40ba38f1, + 0x1b913858, 0xba152f58, 0x3a7f3642, 0x387a39d4, 0x3efcb9af, 0xbcc6c007, 0x3938ba9b, 0xaa27b635, + 0x3509a2e3, 0xbb66b38b, 0xb775bbc3, 0x33f5bbc7, 0xb6bac014, 0x340ba91a, 0x3dadc0ad, 0xb70ea94d, + 0xbcf0b9b4, 0xb38a2f73, 0x304e34f1, 0xbfe63cdf, 0x34c3c329, 0xbdc03c80, 0x39ed37bd, 0xaaa3393d, + 0xbd493a94, 0xb4d2c288, 0xbfb62c3e, 0xb93dbca0, 0x3200b9d6, 0x3875ac36, 0x33f73072, 0xb8f73e47, + 0x3116ba54, 0xc0373a12, 0xb48c3ad7, 0xb798323b, 0xbf50c11c, 0xaf0baf13, 0x2efc3c1c, 0xbc683b04, + 0x38fab611, 0x3f8cb74f, 0xbf1a3b22, 0x3cefc394, 0x26143ca6, 0xbe0337c3, 0xc0723cd0, 0xc15f3f2e, + 0x34da3c11, 0x3cf5b4e7, 0x3030c030, 0x3837b32a, 0xbca2c1c0, 0x3ef133dd, 0xac37982a, 0xb7b33c5e, + 0xb9cdb5b6, 0x2c0cb9e9, 0xab2a3832, 0x419aacea, 0x39223458, 0x3c1fc18e, 0xb7622ae8, 0xbe8c393d, + 0xbd47b009, 0x3abd3d9b, 0x3c56b95a, 0x31073cfe, 0xbc5637d9, 0xb846312d, 0x3d2e3752, 0x403041d2, + 0x3bf0b2d1, 0xbc3930c5, 0xb8aab1ee, 0x3d48405a, 0x2a023c11, 0xc184b800, 0x308d3dd2, 0xbd48b975, + 0xb94232b5, 0xbe0e38c6, 0xc06f3688, 0xbc6133b3, 0xc0943910, 0x38f83cb1, 0xc2b8be18, 0xb3de37d6, + 0x3b2eab85, 0x299cb8f6, 0xbdb13676, 0x3f26b87d, 0xc1173068, 0x3a16b7ae, 0x3b4430ea, 0x30a9b347, + 0x39c331a2, 0x302136bf, 0xbc1338f6, 0x338db882, 0x3bf4b892, 0xba38385d, 0x305d3b4e, 0x387b43aa, + 0xb0c436a4, 0xb925b800, 0xbdefb9d7, 0x4244bba6, 0x345236da, 0xc1913b75, 0x318b381f, 0x405a38a3, + 0x3d74ba2a, 0x3666b9b5, 0x38023990, 0x37ccb8ed, 0x36e93697, 0xb3ef3d43, 0x3cbfbaa8, 0xadc4bdf3, + 0xb0b62c73, 0x386abbd5, 0xbecebaaf, 0x3d5835f2, 0x3ca3358e, 0x3d37b987, 0x3cba3770, 0x400f3d8a, + 0x391037a5, 0x31483bce, 0xb1f2b10b, 0x3a6d343f, 0xb733399b, 0x355038b6, 0xb90abe87, 0xbe4bc1ea, + 0xbe8fbe96, 0xb92a3999, 0xb1c83a77, 0x3508366c, 0xbcc4b7aa, 0xb591b03d, 0xb9fcbecf, 0x40f5bb2f, + 0xc0db4073, 0xbf44421c, 0x34cd339f, 0xb4dc37bd, 0xb44cb531, 0x3b383e9d, 0xa9ffba12, 0xb9ce3655, + 0x3441b32e, 0x3c083f6f, 0xb8603a43, 0x3dfebcf7, 0xbc86c07a, 0xbbf63745, 0x36a62987, 0x267f35e0, + 0xb9f1b4b8, 0x3a39b865, 0xbd262f53, 0x388c37b5, 0xb8893b0e, 0xa4223c18, 0x1666b894, 0xaf72b37f, + 0xb0813772, 0xb0eac090, 0x41d033ad, 0x345fb75b, 0xba823de3, 0xb979afd1, 0x36ec3f52, 0x42c7b928, + 0xb54ebc71, 0x356d3b45, 0x3d77b63d, 0x2873a722, 0xac223ab5, 0x3ccbbe38, 0xba9db82f, 0x34f9b88a, + 0xb9973dde, 0x399dc1b1, 0xc19a3d55, 0x3e72b380, 0xbd0ab701, 0xba7fc16c, 0x3df039a8, 0xb595bf04, + 0xb904ba46, 0x35c93ab7, 0xb9dab181, 0x3a40ba29, 0xb257bee1, 0xbabbc262, 0x39fa3d35, 0x39222661, + 0x3c35b5da, 0xbc02ac1e, 0x375bb9c2, 0xb5e8343c, 0xbd4abc03, 0x38402c70, 0x3a48b98c, 0xb987bd06, + 0x3a3fc008, 0xadf03676, 0x437f391e, 0x38df4179, 0x3dd23a92, 0xb870b42f, 0xb0b33e99, 0x3c3b3136, + 0x3b57b98b, 0x3d80af33, 0x3475b578, 0xb6203f58, 0x386c3a60, 0xbc68af9c, 0xb99a3826, 0xb414b8c8, + 0x374bb952, 0xba0c40e7, 0x40723a49, 0xbc2ab76d, 0xaebbbc67, 0xba0d3a0b, 0x3e433f56, 0xbc433d02, + 0x3e82b9a2, 0xb7963a1f, 0xbc5c355d, 0xba8a3a70, 0xbd61333c, 0xb9c73caf, 0x31f0bd93, 0xbce5bd05, + 0xba81b4b9, 0xbbf2bc7d, 0xb9763740, 0xbe7bb53f, 0xb266bced, 0x3bda30be, 0xbab03f03, 0xbfc73bdf, + 0xbf88adc6, 0x2c42b00b, 0xc2c93e4d, 0x4114be94, 0xbcba3b5a, 0xbd313d7c, 0x37ac3726, 0x3a3a397e, + 0xbcd02f06, 0x3cb93d0c, 0x3f3735a6, 0xb16141f6, 0x3e57ba84, 0x3e1eb423, 0x2e083eb8, 0xbc7fb1c0, + 0xc03a3bef, 0xbb16bcc4, 0x3c4a38a5, 0x3bf83b7f, 0x296439c9, 0xb944bddc, 0x3717406a, 0xaef2b80f, + 0xbf1438c2, 0x37eb3633, 0x3d192c27, 0x37d2ba23, 0x382331dd, 0x39a3bb12, 0xb57cb985, 0x36a5b7fa, + 0x3b5bb882, 0x34c03856, 0xb60d3a69, 0x3b6ec03c, 0x219c3b95, 0xbdd6ac3a, 0x356f4008, 0x33cbae4c, + 0x3c41bd84, 0x3b3c3ba0, 0x387ebb56, 0xba143b22, 0x3f8d33cb, 0x18e03508, 0xa89e3b2d, 0x3c7f3c83, + 0xb82a3cd9, 0x391239b5, 0xbb1c39d4, 0xb1deb8ed, 0xbb923b0b, 0x41783be0, 0x31e83716, 0xbc80b3ff, + 0x393eba07, 0xba8739a9, 0x4185bb2d, 0x32fdb616, 0xbbe7b9d6, 0x3cee3645, 0xbcbab4c6, 0xac743e24, + 0xb9ba2b37, 0x31423c42, 0x33e93505, 0xb53ebe07, 0x3b1db609, 0x36e4bd2b, 0x36fb3c8a, 0xb8c339e9, + 0x42953fba, 0xbf1a3766, 0x3ca339e3, 0xb4343fc7, 0xbb95b899, 0xb2e3398e, 0xb5c8bcd6, 0xb44238c5, + 0x407d3cf2, 0xb847b6e2, 0xbdcd3a6a, 0xb34234b0, 0x2d943a83, 0xb6443e96, 0x3e5dbead, 0x39cdb0bb, + 0xba58a92f, 0xbdf93c0b, 0xa45939ad, 0xbd5c3e21, 0x41c83d83, 0xbeb73299, 0xbc3ab260, 0xb98cb6c7, + 0x3d96321a, 0x3a6f3834, 0x3d353920, 0xbe2828a0, 0xb9dbbe36, 0xbbbab7d3, 0x357fbbd1, 0xc114bd91, + 0xae9bb5fd, 0xb85bc078, 0xb4c3c123, 0xb79638b3, 0x3ab8bb1b, 0x37e533b3, 0x3a311a62, 0x3a67b4f1, + 0x3e7bbbce, 0xbc7cb73c, 0x4098408e, 0x342bbe64, 0x39c73da6, 0x1dea4077, 0x3c46af4e, 0xbbe83ebd, + 0x3cd83936, 0x3e99402f, 0xb2d23c05, 0xb1ebbc89, 0x3d842dfa, 0xbb80bed9, 0xbb7238a5, 0xb173bc2c, + 0xb35db84d, 0x3e9fb0ff, 0x3b57bc79, 0x307dbb5b, 0x3703b498, 0x38bebe4a, 0xab2eb878, 0xbacc3b46, + 0x32d83ab4, 0x365bc20b, 0x36ed35ac, 0x2d1b40e5, 0xb0f7b594, 0xbe4cba7a, 0xb070a99c, 0x2dcab21f, + 0x3fd8bf8a, 0x391bb83e, 0x330fbc35, 0xb69834e7, 0x34c1b482, 0xb81836cd, 0xb4493e69, 0x31ab3646, + 0x341e38c2, 0xb9873925, 0x41b938e5, 0xb1acb79d, 0x38f936c6, 0xa977c103, 0xbce0bead, 0xb897bc63, + 0xb8423f5d, 0x38753d63, 0xad0bab3a, 0xbcd625c0, 0xb85138d6, 0xc077b9a9, 0xb4a13be3, 0x39c4c293, + 0x414b32dc, 0x38cc3ada, 0x3b60bad6, 0xb52e3a85, 0x3d6d3d6c, 0xb855b5a4, 0x27f62ccf, 0xbcf635b1, + 0xb4a034df, 0xb62abe85, 0xad0d39d4, 0xab88b542, 0x345a35dd, 0x3951353b, 0x2f10b8f8, 0x363fbaae, + 0x3ca23878, 0x38da35df, 0xbcfdb443, 0x2fdbb87e, 0x36abbc75, 0x3243389a, 0xba7029dc, 0x32502ee6, + 0x3d29b748, 0xbe2140a9, 0xb60dbba6, 0x3c40bbbc, 0xbb7bb5dd, 0x3bf32631, 0x2ad5bfbb, 0xb41735b7, + 0x3831afde, 0x3145b651, 0xb23f3cfa, 0xb56f39ff, 0x3f03bd12, 0x362333d8, 0x3e043401, 0x3f6bbe9f, + 0xa9d730db, 0xb34d31d9, 0x32529f5e, 0x3c75c0f6, 0xb8ec348d, 0x384138d0, 0x20a83abf, 0x2b97c090, + 0xac01bf35, 0x3722b854, 0x344dc074, 0x39723a91, 0x39b9be82, 0xc02ab62c, 0xba61bb3a, 0x3a292fcb, + 0xbb4eba51, 0x3d17b604, 0x3b16363e, 0x3ab93c14, 0x3a9cb80d, 0x3936b0df, 0xb8cab50c, 0x35f3bd85, + 0xb02b2ffc, 0x38fec0b9, 0xba48390d, 0xbc03408f, 0x3646b937, 0xc0c1c22b, 0xbad0bab4, 0x99a8b3f6, + 0xbf40bc8c, 0xbc802074, 0xb89bb927, 0xa9e533be, 0x3229417f, 0x353cb346, 0xb883395d, 0xbadd3d69, + 0xa7433c6c, 0x3466bd97, 0x3977bf85, 0x39763618, 0xb9f83e77, 0xc024be06, 0x325abc84, 0x3d18bb6c, + 0xbca2b930, 0xbfb2b806, 0xb2533578, 0x33f1b48c, 0xbb21b9ac, 0x374ab012, 0x3c0bb0cc, 0x3dc8aca9, + 0x322cb4d9, 0xb993ab2c, 0x3a913aa6, 0x3da642a1, 0x3e593151, 0x2d8e3751, 0x2d9a3365, 0xb77fbaaa, + 0x32eeb7f0, 0xbb783dd2, 0xbe1c3d0e, 0xbd05be19, 0xbebdbcfe, 0x36b0bba1, 0x40063962, 0x370d380e, + 0x3b9c3e08, 0xb2a82bfe, 0xc0383dd4, 0xb32cc00f, 0x39f333b5, 0xb8933691, 0xac1fbb79, 0xbd823d09, + 0xb867afcd, 0xbbb5c01c, 0xb323bbca, 0x3d5db8d8, 0x2db4bbda, 0x38f9b8ca, 0x3a87b615, 0xc099321c, + 0x3f8c3a26, 0x36f130a7, 0xc2a5386f, 0xb157b485, 0x3f94b4d2, 0xb1bb2056, 0x35dbbb33, 0xba5135b5, + 0x845ebe3f, 0xbd67b7cd, 0xb6702dae, 0xc11f399d, 0x378e2c99, 0xbe92b735, 0x3c614021, 0x34c2ba18, + 0xb927b993, 0xb82ab98a, 0xbf68ba9d, 0xb7f9b507, 0xb71840e7, 0x3cd31e42, 0x3a822b48, 0xbf983ef1, + 0x3a56b0ad, 0xbc1234a6, 0x31d63fe4, 0xa3153d63, 0xb81a22f5, 0x32d73b45, 0x36893e26, 0xbc0c3d81, + 0xbcd33d04, 0xb5c836b2, 0xb86b37bb, 0x3c81b658, 0xba20b81b, 0xadad356a, 0x3048be5d, 0x301db9d4, + 0x3d2dadb0, 0xb415aa8b, 0x413fb4b2, 0xb5163c3c, 0x3572aff4, 0xbbf6b682, 0x3bd3bbce, 0x3d0a3a19, + 0x35923729, 0x3a21bc5f, 0xb6673c5b, 0x38c73c4b, 0xaecd3180, 0xb9d9bd82, 0x36d1ac5a, 0x2c34c22c, + 0xb450bcae, 0x388f3f37, 0xb5fdb698, 0x3b4a23d2, 0xae67b721, 0xab023cbc, 0xbd7339d8, 0xbaf198eb, + 0x381e3ad1, 0xbc12b864, 0xc084c0e9, 0x39713831, 0xb2fa3ca9, 0x3708a167, 0xb9ca405f, 0x31ad3ca1, + 0x339d383f, 0x364f389a, 0x3db4b77d, 0xb2bcb56a, 0xb4ddba8f, 0x32b6bf46, 0xb90a381e, 0xbf8a3db5, + 0x37b4bc4e, 0xad30bc0e, 0x3d833c84, 0x3a043a52, 0x4027bc64, 0x3e6d3c18, 0xaf48bcce, 0x325cb052, + 0xbf1db876, 0xaf28bafa, 0x34983a5b, 0xbc54aab7, 0x3d7e3fdb, 0xbf2039e0, 0xaddcbfc8, 0xb83bb815, + 0x3798b793, 0xb6c53915, 0x395bbc75, 0xb960358f, 0xbcedbc4a, 0x34c2bb4c, 0x20223cfd, 0x30cf2cce, + 0xbc4db632, 0x33083a35, 0x38d22f07, 0x3c45bc58, 0xb9adbae1, 0x3910bb37, 0xb91ebf4b, 0xb27e374d, + 0x367bb59a, 0x39fb3af6, 0x3a2abb7c, 0x2d37bf2d, 0x3803b489, 0xbdc63fc6, 0x3fe4381d, 0xb50833f7, + 0xbeb94023, 0xb0c63d4c, 0x38c5b9aa, 0x360ab65d, 0x3ad7bd31, 0x39bf3f3e, 0xb920bc24, 0xb7a33774, + 0xbcb9afd2, 0x387ab627, 0x38cbb7cc, 0xb37eaf26, 0xbad1b2c5, 0xb982b881, 0xbf5fbe79, 0x3cb6b58c, + 0xc084b76a, 0x305f3888, 0xbb1c4386, 0x2990318f, 0xb415bc45, 0x3b9abd45, 0xb6c7400b, 0x3f4940d8, + 0x39533b3e, 0x41a54121, 0x3a863e15, 0xb50cb745, 0xbc62394d, 0x36fd3ce8, 0xb7d3b3fe, 0x3663320d, + 0x39e838fd, 0x38c0b2dd, 0xb5163992, 0xbc1fb019, 0xa8c635e3, 0x39753fda, 0xc0dd3dba, 0x3863166a, + 0x3ab2a3cd, 0xbcb9332f, 0xb88ebce3, 0x41f6bd5e, 0xbd0aba72, 0xc170bac9, 0xb615bc26, 0xbca7bab0, + 0xbb6cc082, 0x415e3f5b, 0xbe9b3be8, 0x3f323c55, 0x410aa521, 0xbe1ebf27, 0x39dfbca8, 0xbbf5b924, + 0x39c7b2aa, 0x424b3af7, 0xaed0b4c6, 0x3532358e, 0x39243116, 0xc199bee5, 0x3849ba6e, 0xb9203c77, + 0xb4813c82, 0x3c4e33de, 0x3775ba98, 0x96a7b594, 0x4000bd85, 0x3fd82fd0, 0x3bb73d71, 0x38083dc5, + 0x34a8bc1c, 0xb86c3d3d, 0x38413bf1, 0xb9d3399c, 0xb752bcee, 0xb00a3ab9, 0x3198b744, 0x398f2fcd, + 0x3d759a32, 0x3bbc360a, 0xbb5937e9, 0xb330b969, 0x3ae6c0a5, 0x40b03e0c, 0x343cbcda, 0xb86c34bc, + 0x35be3cb0, 0xac1bbd82, 0xa5a7bb08, 0xb0f534bf, 0xb4853087, 0xbb9ca123, 0xb45d3df9, 0x404c3c17, + 0x39a3347a, 0xb913b6e8, 0x38cebd20, 0x40b0c006, 0xbb4d402a, 0x3e493e7d, 0xb6c83b8b, 0x3b902d82, + 0xae29388d, 0x3ad7327b, 0xb9843d10, 0x34cc3933, 0xbbda3002, 0x3819bcf3, 0xbc312c89, 0x38e83dbd, + 0xbb87b4e4, 0x3886b172, 0xb82939b9, 0xb345bd54, 0xba7fb951, 0x36c13d9f, 0x3dd3b259, 0x36a43a49, + 0x3cf23b65, 0xbe2cacae, 0x3ba3bdae, 0x3603356a, 0x33703d0e, 0x2d88b96b, 0x36e3b8d3, 0x3b24b804, + 0xbdd0c000, 0xb8c7c006, 0x33673dc9, 0x344bb663, 0x341f3406, 0xc0cfac64, 0x3d0234c6, 0x39f83676, + 0x2e04a4b0, 0x3b5fb4ee, 0x37d4bb6d, 0xb491b0cc, 0xacc8b1aa, 0x3c6cbea5, 0xbc46b99e, 0x3b6cbe17, + 0xbf93bd5d, 0x38fab462, 0x331fc0e4, 0x31643419, 0x4066bb25, 0xbe28bd61, 0xbc663891, 0xbc0c32be, + 0x3dee3055, 0xb3ab20f7, 0x3a5031ff, 0xbad6be6a, 0x3cb5bb8f, 0x3e97be0a, 0xbc8b3110, 0x3c1b36cb, + 0xbe5fbb83, 0x310baecd, 0x2985b4b8, 0xb9d93c4a, 0xbc1d39c5, 0x3d9e390b, 0xbca43675, 0xb49c3ce1, + 0x3fffbb08, 0x35e33927, 0xbc60bd50, 0xbd1339b6, 0xc2a24123, 0x37d9beab, 0x3425b8ed, 0x3a36bace, + 0x4057b0c7, 0x3e0c3c69, 0x3a03b5b0, 0xbaa0ba05, 0x9c4c19c4, 0x3c37b72f, 0xbd023467, 0xb537b522, + 0x2c4d3673, 0xbdefb17b, 0x3daab95a, 0xb4e3bbe8, 0xbb3a3034, 0xbd9cc116, 0xc1d53891, 0xb874bd9f, + 0x357fbc4b, 0x3f79b09a, 0xc07e3f0c, 0xbca2bbfe, 0x2af5bb50, 0xbaf7b02f, 0xb509b629, 0x36fc39c8, + 0x395ca825, 0x3553b901, 0xbbebbd9d, 0xb45139f2, 0x3f3c37bc, 0x386635a3, 0x35fb3046, 0xbcb33999, + 0x3d1da417, 0x3becbc37, 0x3bfa3c7c, 0xc1d5bcb8, 0x3c403c1a, 0xb73bb23c, 0x3a093546, 0xbb7fbbf5, + 0xbe3f30e8, 0x3a8b39fe, 0x9b02341a, 0x2b91bc35, 0x3a514256, 0xba433fd1, 0x393ab55c, 0xae6a3c6f, + 0xb2663da8, 0xba78bd17, 0x3e3a35cf, 0x3e4db5dd, 0x32ff345d, 0x35a2bca2, 0x3c82b9f1, 0x38f1bbc5, + 0xbe7b3f36, 0xb44abe77, 0x3501bd3b, 0xbe42344c, 0x3dc3bc81, 0xbf152613, 0x3c71bb1c, 0xba6a345f, + 0x3c03b944, 0xbae1b5ea, 0x3e77ab67, 0xb9dbc04f, 0xb86ebe0c, 0xb9cdbc4f, 0xb07a3ca8, 0xb41f36be, + 0x3256ac27, 0xb3273870, 0x3b9db5f6, 0xb8de30ba, 0xb06a3c1b, 0x3d543ae1, 0x3d08ba28, 0xb93838a3, + 0xb96b3472, 0x3b803945, 0x3ab1adbe, 0x39b53ccf, 0x383634a8, 0x308cb49e, 0xb5ba3916, 0x3b1bbb8a, + 0xbe6fbc56, 0xb4d5bd9e, 0xa3433e83, 0xb2fc3adc, 0x3b55ba2f, 0xb6243cec, 0x37d7261f, 0x3eb1bf76, + 0xb9c838b1, 0x361ca60d, 0xc00cbca6, 0xbc6bbb55, 0x38ba34c5, 0x383cb802, 0x3e68bd59, 0xbbad3a2e, + 0x3b243c74, 0x3ebf3010, 0xbfa8b953, 0x3b7cb16b, 0x2c823952, 0x412cb8ef, 0x3cdbb7ec, 0x38e03bc9, + 0xbaa3bbfb, 0xb1df3f28, 0xbccf3a77, 0xb933b1a5, 0x37cb3018, 0x31bcbec0, 0x3997bd2e, 0x39063930, + 0x37fdb86f, 0x33e5323d, 0x38d7a60a, 0xb451b30a, 0x31233597, 0xbe8ac063, 0x37afbe42, 0x3b21aff0, + 0xb277365f, 0x2ad2b5b8, 0x3a45c11d, 0xb50031fd, 0xbc72bf3f, 0x24a33038, 0xb5952890, 0xc1c7c08f, + 0xbcaab86e, 0x38a63189, 0x325db408, 0x3b052ca7, 0xbf353c55, 0x3b7e387c, 0x3a0b3c20, 0x3db83741, + 0x3c4d3a12, 0x3d11b87c, 0xbed7bdc9, 0xbdbb2444, 0x3cc9b9c6, 0xba10ae72, 0x395338b7, 0xbbe2b1e6, + 0x30b1a918, 0x283dbc8f, 0xb160aeb3, 0x33363794, 0x401d3327, 0x3d15bf99, 0x3a32ad02, 0x3c573df6, + 0x34edbef4, 0x368cbc1f, 0xb95b3aed, 0x3bf338dd, 0x3c2f4185, 0x37b5bbfc, 0x355e37f3, 0x3c463559, + 0x38bcc003, 0xbda03b53, 0x2acbac8d, 0x396cb8ac, 0xbb1f2be0, 0x3b903ce8, 0x3ef8ae9f, 0x3cf4b33e, + 0xb8982c97, 0x3f56baa5, 0xbb163e58, 0xc004b33c, 0xbc68400e, 0xa93938bd, 0x3b353a55, 0x3db540d6, + 0xb8694176, 0xbd7f30ba, 0xb379ba99, 0xb9153ecf, 0x30f2b19f, 0xb6323e11, 0x3be5be94, 0xac5d3d28, + 0xbb64b6e7, 0x4380b282, 0x36b5bde1, 0x386736c2, 0x2e1c3cca, 0x2e2f31cc, 0xbcffb0d4, 0xbd67bb1c, + 0x3d2c309a, 0xbd03ab7d, 0xbefcbbe1, 0x3b98bfbe, 0xbc72b53c, 0x3cca2f15, 0xbc36c106, 0x3b57b702, + 0xb4da391a, 0x3afe3861, 0x335b3ea8, 0xbd6234e4, 0xbc863f1f, 0x35bcbcf4, 0xbf53b34a, 0xbd103834, + 0x39a1273d, 0xb4e4b839, 0x34da39b8, 0x3a76bed6, 0xc0fa3c50, 0xb061bd35, 0x34f0b1b6, 0xbc183ab1, + 0xb62cb91d, 0xb829bb36, 0xafabb940, 0x3e813d46, 0xb822b20b, 0x35fdbe39, 0x396632f1, 0xb5432e17, + 0x3a89bbed, 0xbcc5b2a3, 0xb42b3369, 0x1f58bbd4, 0x3ef4bc48, 0xac2b3bbe, 0xb2e4bdf5, 0x3f94bcae, + 0xc068a11a, 0xb82e35bb, 0xbbb93f77, 0x3b46b386, 0x22f3b675, 0xbcb23e8b, 0xba29b8ba, 0xbd084235, + 0x3f9f1038, 0x399fade8, 0xb9da3363, 0x37c83909, 0xbc4c3c82, 0xb788bc3a, 0x39daa897, 0x3beeb6d2, +}; +// END tq_lut +// the trellis codebook point i (two fp16 values in one word): from the constant copy, or the threadgroup one +static inline uint g_tq(threadgroup const uint* tab, uint i) { return ctab ? tq_lut_u[i] : tab[i]; } + +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template static inline void load_tab_tq(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 2048; i += nth) tab[i] = tq_lut_u[i]; +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { load_tab_tq(tab, tid, nth); } +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { load_tab_tq(tab, tid, nth); } +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { load_tab_tq(tab, tid, nth); } +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} + +static inline void load_signs(threadgroup uint* tab, int at, ushort tid, ushort nth) { + threadgroup uchar* s = (threadgroup uchar*)(tab + at); + for (int i = tid; i < 128; i += nth) s[i] = ksigns_iq2xs[i]; +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 256; i += nth) tab[i] = iq3xxs_grid[i]; + load_signs(tab, 256, tid, nth); +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 512; i += nth) tab[i] = iq3s_grid[i]; + load_signs(tab, 512, tid, nth); +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 256; i += nth) { tab[2 * i] = uint(iq2xxs_grid[i]); tab[2 * i + 1] = uint(iq2xxs_grid[i] >> 32); } + load_signs(tab, 512, tid, nth); +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 512; i += nth) { tab[2 * i] = uint(iq2xs_grid[i]); tab[2 * i + 1] = uint(iq2xs_grid[i] >> 32); } + load_signs(tab, 1024, tid, nth); +} +// IQ1_M: no threadgroup table - the tile reads iq1s_grid from the constants (the w4 helper in moe_mm.metal) +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} +template <> inline void load_tab(threadgroup uint*, ushort, ushort) {} // (same: iq1s_grid from the constants) +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + if (ctab) return; + for (int i = tid; i < 1024; i += nth) { tab[2 * i] = uint(iq2s_grid[i]); tab[2 * i + 1] = uint(iq2s_grid[i] >> 32); } + load_signs(tab, 2048, tid, nth); +} +// IQ4: a byte -> (codebook[low nibble], codebook[high nibble]) as half2, so one lookup gives two weights +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + for (int i = tid; i < 256; i += nth) + tab[i] = as_type(half2(half(kvalues_iq4nl[i & 15]), half(kvalues_iq4nl[i >> 4]))); +} +template <> inline void load_tab(threadgroup uint* tab, ushort tid, ushort nth) { + load_tab(tab, tid, nth); +} + +// ---- unit decoders: return s, fill q[UNIT] ------------------------------------------------------------------------ +// ui = unit index inside the block (0 .. QK/UNIT - 1). +template static inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab); + +// Q2_0: 64 = 4 units of 16; code c (2 bits, LSB first) -> c - 1; scale d. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2 + 4 * ui); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float((a >> (2 * j)) & 3) - 1.0f; + q[8 + j] = float((c >> (2 * j)) & 3) - 1.0f; + } + return rdh(b); +} + +// 8 signed grid bytes from a 4-byte grid word pair helper +static inline void put_signed4(thread float* q, uint g, uint s) { + const uchar4 v = as_type(g); + q[0] = (s & 1) ? -float(v[0]) : float(v[0]); + q[1] = (s & 2) ? -float(v[1]) : float(v[1]); + q[2] = (s & 4) ? -float(v[2]) : float(v[2]); + q[3] = (s & 8) ? -float(v[3]) : float(v[3]); +} + +// IQ3_XXS: 256 = 8 x 32; unit = half of a 32 (two grid pairs). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint aux = rd32(b + 2 + 64 + 4 * ib); + device const uchar* qs = b + 2 + 8 * ib + 4 * h; + for (int l2 = 0; l2 < 2; ++l2) { + const int l = 2 * h + l2; + const uint s = ksg(tab, 256, (aux >> (7 * l)) & 127); + put_signed4(q + 8 * l2, g_iq3xxs(tab, qs[2 * l2]), s); + put_signed4(q + 8 * l2 + 4, g_iq3xxs(tab, qs[2 * l2 + 1]), s >> 4); + } + return rdh(b) * (0.5f + float(aux >> 28)) * 0.5f; +} + +// IQ2_XXS: per 32: two u32 (4 grid indices, then 4 x 7 sign bits + 4-bit scale). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint a0 = rd32(b + 2 + 8 * ib), a1 = rd32(b + 2 + 8 * ib + 4); + for (int l2 = 0; l2 < 2; ++l2) { + const int l = 2 * h + l2; + const uint idx = (a0 >> (8 * l)) & 255; + const uint s = ksg(tab, 512, (a1 >> (7 * l)) & 127); + put_signed4(q + 8 * l2, g_iq2xxs(tab, idx, 0), s); + put_signed4(q + 8 * l2 + 4, g_iq2xxs(tab, idx, 1), s >> 4); + } + return rdh(b) * (0.5f + float(a1 >> 28)) * 0.25f; +} + +// IQ2_XS: uint16 qs[32] (9-bit grid index + 7-bit sign index), scales[8] nibbles. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + device const ushort* qs = (device const ushort*)(b + 2); + for (int l2 = 0; l2 < 2; ++l2) { + const uint v = qs[4 * ib + 2 * h + l2]; + const uint idx = v & 511, s = ksg(tab, 1024, v >> 9); + put_signed4(q + 8 * l2, g_iq2xs(tab, idx, 0), s); + put_signed4(q + 8 * l2 + 4, g_iq2xs(tab, idx, 1), s >> 4); + } + const uint sc = b[66 + ib]; + return rdh(b) * (0.5f + float((sc >> (4 * h)) & 15)) * 0.25f; +} + +// IQ2_S: qs[32] low 8 bits of the 10-bit index, signs[32], qh[8] (2 high bits per group of 8), scales[8]. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint qh = b[66 + ib]; + for (int l2 = 0; l2 < 2; ++l2) { + const int l = 2 * h + l2; + const uint idx = uint(b[2 + 4 * ib + l]) | ((qh << (8 - 2 * l)) & 0x300); + const uint s = b[2 + 32 + 4 * ib + l]; + put_signed4(q + 8 * l2, g_iq2s(tab, idx, 0), s); + put_signed4(q + 8 * l2 + 4, g_iq2s(tab, idx, 1), s >> 4); + } + const uint sc = b[74 + ib]; + return rdh(b) * (0.5f + float((sc >> (4 * h)) & 15)) * 0.25f; +} + +// IQ3_S: qs[64] low 8 bits of 9-bit indices, qh[8], signs[32], scales[4] (nibble per 32). +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + const int ib = ui >> 1, h = ui & 1; + const uint qh = b[66 + ib]; + for (int l2 = 0; l2 < 2; ++l2) { + const int l = 2 * h + l2; + const uint i1 = uint(b[2 + 8 * ib + 2 * l]) | ((qh << (8 - 2 * l)) & 256); + const uint i2 = uint(b[2 + 8 * ib + 2 * l + 1]) | ((qh << (7 - 2 * l)) & 256); + const uint s = b[74 + 4 * ib + l]; + put_signed4(q + 8 * l2, g_iq3s(tab, i1), s); + put_signed4(q + 8 * l2 + 4, g_iq3s(tab, i2), s >> 4); + } + const uint sc = b[106 + (ib >> 1)]; + return rdh(b) * float(1 + 2 * ((sc >> (4 * (ib & 1))) & 15)); +} + +// IQ4_NL: 32 = 1 unit; low nibbles are elements 0..15, high 16..31 (one half2 lookup per byte). +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint* tab) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + q[2 * j] = float(w0.x); + q[2 * j + 16] = float(w0.y); + q[2 * j + 1] = float(w1.x); + q[2 * j + 17] = float(w1.y); + } + return rdh(b); +} + +// IQ4_XS: 256 = 8 units of 32; 6-bit scale from scales_l nibble + 2 bits of scales_h, minus 32. +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + device const ushort* p = (device const ushort*)(b + 8 + 16 * ui); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const half2 w0 = as_type(tab[v & 255]), w1 = as_type(tab[v >> 8]); + q[2 * j] = float(w0.x); + q[2 * j + 16] = float(w0.y); + q[2 * j + 1] = float(w1.x); + q[2 * j + 17] = float(w1.y); + } + const uint sh = *(device const ushort*)(b + 2); + const int ls = int((b[4 + (ui >> 1)] >> (4 * (ui & 1))) & 15) | int(((sh >> (2 * ui)) & 3) << 4); + return rdh(b) * float(ls - 32); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 16; ++j) { + const ushort v = p[j]; + q[2 * j] = float(as_type(uchar(v & 255))); + q[2 * j + 1] = float(as_type(uchar(v >> 8))); + } + return rdh(b); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + device const ushort* p = (device const ushort*)(b + 2); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + q[2 * j] = float(v & 15) - 8.0f; + q[2 * j + 16] = float((v >> 4) & 15) - 8.0f; + q[2 * j + 1] = float((v >> 8) & 15) - 8.0f; + q[2 * j + 17] = float(v >> 12) - 8.0f; + } + return rdh(b); +} + +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + const uint hb = rd32(b + 2); + device const ushort* p = (device const ushort*)(b + 6); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const int j0 = 2 * j, j1 = 2 * j + 1; + q[j0] = float(int(v & 15) | int(((hb >> j0) & 1) << 4)) - 16.0f; + q[j0 + 16] = float(int((v >> 4) & 15) | int(((hb >> (j0 + 16)) & 1) << 4)) - 16.0f; + q[j1] = float(int((v >> 8) & 15) | int(((hb >> j1) & 1) << 4)) - 16.0f; + q[j1 + 16] = float(int(v >> 12) | int(((hb >> (j1 + 16)) & 1) << 4)) - 16.0f; + } + return rdh(b); +} + +// Q5_1: 32 = 1 unit; q = the 5-bit code (no offset; the min m is applied through the x sums), returns d +template <> inline float deq(device const uchar* b, int, thread float* q, threadgroup const uint*) { + const uint hb = rd32(b + 4); + device const ushort* p = (device const ushort*)(b + 8); + for (int j = 0; j < 8; ++j) { + const ushort v = p[j]; + const int j0 = 2 * j, j1 = 2 * j + 1; + q[j0] = float(int(v & 15) | int(((hb >> j0) & 1) << 4)); + q[j0 + 16] = float(int((v >> 4) & 15) | int(((hb >> (j0 + 16)) & 1) << 4)); + q[j1] = float(int((v >> 8) & 15) | int(((hb >> j1) & 1) << 4)); + q[j1 + 16] = float(int(v >> 12) | int(((hb >> (j1 + 16)) & 1) << 4)); + } + return rdh(b); +} + +// Trellis types (ggml-common-tq.h): a 128-value block is {fp16 d; qs}, 32 steps of K bits (K = 8 / 6 / 7). A unit is +// lane ui = steps 4ui..4ui+3, 16 weights. Step t's 15-bit state: TQ2_T qs[(t + 31) % 32] << 8 | qs[t]; TQK the 16 +// stream bits from bit (31 - t) * K, circular over 32 K bits. x = state * 0x9e3779b1 selects two codebook points, +// lut[x >> 21] and lut[(x >> 10) & 2047] (fp16 pairs); the four weights of the step are those four values times d. +// Same arithmetic as src/kernels/cuda/iq_kernels.cu (tq_states4 / tq_step) and ref_quant.cpp (deq_tq). +template +static inline float deq_tq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { + device const ushort* qs16 = (device const ushort*)(b + 2); + uint s[4]; + if (KB == 8) { + const uint w01 = qs16[2 * ui], w23 = qs16[2 * ui + 1]; + const uint prev = b[2 + ((4 * ui + 31) & 31)]; + const uint b0 = w01 & 0xFFu, b1 = w01 >> 8, b2 = w23 & 0xFFu, b3 = w23 >> 8; + s[0] = (prev << 8) | b0; + s[1] = (b0 << 8) | b1; + s[2] = (b1 << 8) | b2; + s[3] = (b2 << 8) | b3; + } else { + constexpr int NW = 2 * KB; + const int off3 = (28 - 4 * ui) * KB, wb = off3 >> 4, r = off3 & 15; + const int wi1 = wb + 1 >= NW ? wb + 1 - NW : wb + 1; + const int wi2 = wb + 2 >= NW ? wb + 2 - NW : wb + 2; + const int wi3 = wb + 3 >= NW ? wb + 3 - NW : wb + 3; + const ulong v = ulong(qs16[wb]) | (ulong(qs16[wi1]) << 16) | (ulong(qs16[wi2]) << 32) | (ulong(qs16[wi3]) << 48); + for (int i = 0; i < 4; ++i) s[i] = uint(v >> (r + (3 - i) * KB)) & 0xFFFFu; + } + for (int i = 0; i < 4; ++i) { + const uint x = (s[i] & 0x7FFFu) * 0x9e3779b1u; + const float2 a = float2(as_type(g_tq(tab, x >> 21))); + const float2 c = float2(as_type(g_tq(tab, (x >> 10) & 2047u))); + q[4 * i] = a.x; + q[4 * i + 1] = a.y; + q[4 * i + 2] = c.x; + q[4 * i + 3] = c.y; + } + return rdh(b); +} +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { return deq_tq<8>(b, ui, q, tab); } +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { return deq_tq<6>(b, ui, q, tab); } +template <> inline float deq(device const uchar* b, int ui, thread float* q, threadgroup const uint* tab) { return deq_tq<7>(b, ui, q, tab); } + +// sum over groups of LPR consecutive lanes (LPR = 1, 2, 4, 8, 16 or 32); the result is valid in every lane +template +static inline float lane_sum(float v) { + if (LPR == 32) return simd_sum(v); + for (ushort o = LPR / 2; o > 0; o >>= 1) v += simd_shuffle_xor(v, o); + return v; +} + +// ---- Gyro: the block Hadamard rotation of the activations ------------------------------------------------------------ +// Hadamard-folded experts (docs/GYRO.md, include/strata/kernels/cpu/hadamard.hpp) take y = H (s * x) per block of 128 +// values (the normalized Sylvester-Walsh-Hadamard matrix, s = a +/-1 sign vector of the width) in place of x: the gate/up +// input and the down input. One thread per value, 128 threads a block, grid (width / 128, rows). The butterflies run in +// the CPU's order (len = 1, 2, 4, ...; lower element a + c, upper a - c), so the rotated floats are the CPU's and the +// CUDA kernel's: len 1..16 by SIMD shuffles, 32 and 64 through threadgroup memory. +struct HadArgs { + uint width; // values per row (a multiple of 128) + uint x_stride; // row strides in elements + uint y_stride; + float scale; // 1 / sqrt(128), as the CPU computes it + uint has_signs; + uint pad; +}; +template +static inline float had_ld(device const T* p); +template <> inline float had_ld(device const float* p) { return *p; } +template <> inline float had_ld(device const half* p) { return float(*p); } +template +static inline void had_body(device const T* x, device T* y, device const float* signs, constant HadArgs& a, uint3 tg, + ushort tid, threadgroup float* buf) { + const uint col = tg.x * 128u + tid; + float v = had_ld(x + ulong(tg.y) * a.x_stride + col); + if (a.has_signs) v *= signs[col]; + for (ushort len = 1; len < 32; len <<= 1) { + const float o = simd_shuffle_xor(v, len); + v = (tid & len) ? o - v : v + o; + } + buf[tid] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float o = buf[tid ^ 32]; + v = (tid & 32) ? o - v : v + o; + threadgroup_barrier(mem_flags::mem_threadgroup); + buf[tid] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + o = buf[tid ^ 64]; + v = (tid & 64) ? o - v : v + o; + y[ulong(tg.y) * a.y_stride + col] = T(v * a.scale); +} +kernel void had128_f32(device const float* x [[buffer(0)]], device float* y [[buffer(1)]], + device const float* signs [[buffer(2)]], constant HadArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float buf[128]; + had_body(x, y, signs, a, tg, tid, buf); +} +kernel void had128_f16(device const half* x [[buffer(0)]], device half* y [[buffer(1)]], + device const float* signs [[buffer(2)]], constant HadArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float buf[128]; + had_body(x, y, signs, a, tg, tid, buf); +} +// half rows in place, one threadgroup a row (grid.x = the row, so an indirect grid of the row count works): the prompt +// path's h rows (the padded rows of the expert tiles), width values each +kernel void had128_f16_rows(device half* x [[buffer(0)]], device half* y [[buffer(1)]], + device const float* signs [[buffer(2)]], constant HadArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + threadgroup float buf[128]; + for (uint b = 0; b < a.width / 128u; ++b) { + had_body(x, y, signs, a, uint3(b, tg.x, 0), tid, buf); + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} diff --git a/edge0-strata/metal/kernels/router.metal b/edge0-strata/metal/kernels/router.metal new file mode 100644 index 00000000..e5c75274 --- /dev/null +++ b/edge0-strata/metal/kernels/router.metal @@ -0,0 +1,104 @@ +// metal/kernels/router.metal - the router GEMV and moe_route in ONE dispatch (after quant/units/moe/gemv.metal): the +// logits rows as gemv's (LPR = 32 lanes a row, KS = 1: bitwise gemv__32_1_*_1's), 16 rows a threadgroup of 512 +// threads; every threadgroup bumps a device counter after its rows (device-scope fences around it), and the LAST one +// runs moe_route's body (the softmax / top-10 / weights, then the plan) with its 512 threads (= n_expert) and resets +// the counter. No threadgroup waits for another. ids / w / plan are bitwise moe_route's. +// Buffers: W, xp, s16, logits, args (GemvArgs: n_out = n_expert), ids, w, plan, MoeArgs, counter (one uint, zero). + +static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan, + constant MoeArgs& a, ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp, + threadgroup int* s_ent) { + if (sg < a.n_tok) { + device const float* l = logits + sg * 512; + float v[16]; + for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i]; + float mx = -INFINITY; + for (int i = 0; i < 16; ++i) mx = max(mx, v[i]); + for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m)); + float sum = 0.0f; + for (int i = 0; i < 16; ++i) { + v[i] = exp(v[i] - mx); + sum += v[i]; + } + for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m); + const float rcp = 1.0f / sum; + for (int i = 0; i < 16; ++i) { + v[i] *= rcp; + if (isnan(v[i])) v[i] = -FLT_MAX; + } + float selected = 0.0f, ssum = 0.0f; + for (int rank = 0; rank < 10; ++rank) { + float best = v[0]; + int ex = lane; + for (int i = 1; i < 16; ++i) + if (v[i] > best) { best = v[i]; ex = lane + 32 * i; } + for (ushort m = 16; m > 0; m >>= 1) { + const float ob = simd_shuffle_xor(best, m); + const int oe = simd_shuffle_xor(ex, m); + if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; } + } + if ((ex & 31) == lane) { + v[ex / 32] = -INFINITY; + ids[sg * 10 + rank] = ex; + ssum += best; + } + if (rank == lane) selected = best; + } + for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m); + ssum = max(ssum, 6.103515625e-5f); + const float inv = 1.0f / ssum; + if (lane < 10) w[sg * 10 + lane] = selected * inv; + } + threadgroup_barrier(mem_flags::mem_device); + plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); +} + +template +kernel void router_route(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]], + device const float* s16 [[buffer(2)]], device float* logits [[buffer(3)]], + constant GemvArgs& ga [[buffer(4)]], device int* ids [[buffer(5)]], device float* w [[buffer(6)]], + device int* plan [[buffer(7)]], constant MoeArgs& a [[buffer(8)]], + device atomic_uint* counter [[buffer(9)]], + uint tgx [[threadgroup_position_in_grid]], uint ngx [[threadgroups_per_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort nth [[threads_per_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int UNIT = fmt::UNIT; + threadgroup uint tab[fmt::TAB]; + threadgroup int s_grp[32], s_ent[32]; + threadgroup uint last; + load_tab(tab, tid, nth); + threadgroup_barrier(mem_flags::mem_threadgroup); + // the GEMV: simdgroup sg -> row tg.x * 16 + sg, 32 lanes over its units + const int row = int(tgx) * 16 + sg; + const int rr = min(row, int(ga.n_out) - 1); + device const uchar* rp = W + (ulong)rr * ga.row_bytes; + const int U = int(ga.n_in) / UNIT, s16n = int(ga.n_in) / 16; + float acc[8]; + for (int t = 0; t < NTOK; ++t) acc[t] = 0.0f; + for (int u = lane; u < U; u += 32) { + float q[UNIT]; + float4 sc; + Dec::run(rp, u, q, sc, tab); + for (int t = 0; t < NTOK; ++t) { + float xs[UNIT], xsum[2]; + X::load(xp + t * int(ga.x_stride), s16 + t * s16n, u, xs, xsum); + acc[t] += unit_combine(q, sc, xs, xsum); + } + } + for (int t = 0; t < NTOK; ++t) { + const float v = lane_sum<32>(acc[t]); + if (lane == 0 && row < int(ga.n_out)) logits[t * int(ga.y_stride) + row] = v; + } + // the last threadgroup to finish routes + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + threadgroup_barrier(mem_flags::mem_device); + if (tid == 0) last = atomic_fetch_add_explicit(counter, 1u, memory_order_relaxed) == ngx - 1 ? 1u : 0u; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (!last) return; + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + route_body(logits, ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent); + if (tid == 0) atomic_store_explicit(counter, 0u, memory_order_relaxed); +} + +template [[host_name("router_route_30")]] kernel decltype(router_route) router_route; +template [[host_name("router_route_0")]] kernel decltype(router_route) router_route; diff --git a/edge0-strata/metal/kernels/rt_bench.metal b/edge0-strata/metal/kernels/rt_bench.metal new file mode 100644 index 00000000..bdce780a --- /dev/null +++ b/edge0-strata/metal/kernels/rt_bench.metal @@ -0,0 +1,67 @@ +// metal/kernels/rt_bench.metal - bandwidth and dispatch-cost probes (rt_bench.cpp). + +// Grid-stride streaming read: thread i reads src[i], src[i + nt], ... (coalesced across the grid). +kernel void k_read(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 a = 0; + for (uint i = gid; i < n; i += nt) a += src[i]; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// The same, 4 independent loads in flight per iteration. +kernel void k_read_u4(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + uint4 a = 0, b = 0, c = 0, d = 0; + uint i = gid; + for (; i + 3 * nt < n; i += 4 * nt) { + a += src[i]; + b += src[i + nt]; + c += src[i + 2 * nt]; + d += src[i + 3 * nt]; + } + for (; i < n; i += nt) a += src[i]; + a += b + c + d; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// Contiguous chunk per threadgroup (how a GEMV row block reads): threadgroup g reads [g*chunk, (g+1)*chunk). +kernel void k_read_chunk(device const uint4* src [[buffer(0)]], device uint* out [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& chunk [[buffer(3)]], + uint tid [[thread_position_in_threadgroup]], uint tgs [[threads_per_threadgroup]], + uint g [[threadgroup_position_in_grid]], uint gid [[thread_position_in_grid]]) { + uint4 a = 0, b = 0; + const uint base = g * chunk, end = min(base + chunk, n); + uint i = base + tid; + for (; i + tgs < end; i += 2 * tgs) { + a += src[i]; + b += src[i + tgs]; + } + for (; i < end; i += tgs) a += src[i]; + a += b; + out[gid] = a.x ^ a.y ^ a.z ^ a.w; +} + +// Read + write (a state update): dst[i] = src[i] * 1.0001 for float4. +kernel void k_copy(device const float4* src [[buffer(0)]], device float4* dst [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint nt [[threads_per_grid]]) { + for (uint i = gid; i < n; i += nt) dst[i] = src[i] * 1.0001f; +} + +// Dispatch-cost probes. +kernel void k_empty(device uint* out [[buffer(0)]], uint gid [[thread_position_in_grid]]) { + if (gid == 0xFFFFFFFFu) out[0] = 1; +} + +// A dependent chain: each dispatch reads what the previous one wrote. +kernel void k_chain(device uint* buf [[buffer(0)]], uint gid [[thread_position_in_grid]]) { + if (gid == 0) buf[0] = buf[0] + 1; +} + +// A small but real kernel: 2560 floats read, 2560 written (an elementwise op on one token). +kernel void k_small(device const float* x [[buffer(0)]], device float* y [[buffer(1)]], + uint gid [[thread_position_in_grid]]) { + if (gid < 2560) y[gid] = x[gid] * 0.5f + y[gid]; +} diff --git a/edge0-strata/metal/kernels/sample.metal b/edge0-strata/metal/kernels/sample.metal new file mode 100644 index 00000000..8443d360 --- /dev/null +++ b/edge0-strata/metal/kernels/sample.metal @@ -0,0 +1,59 @@ +// metal/kernels/sample.metal - greedy sampling over the LM head's logits: the largest logit, ties to the lowest +// token id (NaN never wins). Two stages: argmax_part (grid (parts, n_tok), 256 threads) writes one (value, id) +// pair per threadgroup; argmax_final (grid (n_tok), 256 threads) reduces them in order. + +struct ArgArgs { + uint n_vocab; + uint parts; +}; + +static inline void better(thread float& bv, thread int& bi, float v, int i) { + if (v > bv || (v == bv && i < bi)) { bv = v; bi = i; } +} + +static inline void tg_argmax(thread float& bv, thread int& bi, threadgroup float* sv, threadgroup int* si, ushort sg, + ushort lane) { + for (ushort m = 16; m > 0; m >>= 1) { + const float ov = simd_shuffle_xor(bv, m); + const int oi = simd_shuffle_xor(bi, m); + better(bv, bi, ov, oi); + } + if (lane == 0) { sv[sg] = bv; si[sg] = bi; } + threadgroup_barrier(mem_flags::mem_threadgroup); + bv = -INFINITY; + bi = 0x7fffffff; + for (int k = 0; k < 8; ++k) better(bv, bi, sv[k], si[k]); +} + +kernel void argmax_part(device const float* logits [[buffer(0)]], device float* pv [[buffer(1)]], + device int* pi [[buffer(2)]], constant ArgArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float sv[8]; + threadgroup int si[8]; + const uint t = tg.y, per = (a.n_vocab + a.parts - 1) / a.parts; + const uint b0 = tg.x * per, b1 = min(b0 + per, a.n_vocab); + device const float* l = logits + t * a.n_vocab; + float bv = -INFINITY; + int bi = 0x7fffffff; + for (uint i = b0 + tid; i < b1; i += 256) { + const float v = l[i]; + if (!isnan(v)) better(bv, bi, v, int(i)); + } + tg_argmax(bv, bi, sv, si, sg, lane); + if (tid == 0) { pv[t * a.parts + tg.x] = bv; pi[t * a.parts + tg.x] = bi; } +} + +kernel void argmax_final(device const float* pv [[buffer(0)]], device const int* pi [[buffer(1)]], + device int* out [[buffer(2)]], constant ArgArgs& a [[buffer(3)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float sv[8]; + threadgroup int si[8]; + const uint t = tg.x; + float bv = -INFINITY; + int bi = 0x7fffffff; + for (uint k = tid; k < a.parts; k += 256) better(bv, bi, pv[t * a.parts + k], pi[t * a.parts + k]); + tg_argmax(bv, bi, sv, si, sg, lane); + if (tid == 0) out[t] = bi; +} diff --git a/edge0-strata/metal/kernels/units.metal b/edge0-strata/metal/kernels/units.metal new file mode 100644 index 00000000..15f584cb --- /dev/null +++ b/edge0-strata/metal/kernels/units.metal @@ -0,0 +1,321 @@ +// metal/kernels/units.metal - the dot of one unit with the window's tokens (after quant.metal). +// +// X::load(xrow, u, xs, xsum) prepares unit u's x slice of one token (xrow = the token's row start), once per unit +// and shared by every weight row the lane reads: the UNIT x values in the order the decoder produces weights, and +// up to two partial sums (for the formats with an offset or a min). Where the codes are read with the mask trick +// (an AND keeps a field in place, so the code arrives multiplied by 2^k), the x values are pre-scaled by 2^-k here: +// exact in float. + +// the threadgroup-memory form (x staged once per threadgroup; the same values) +static inline void load16(threadgroup const float* p, thread float* xs) { + threadgroup const float4* v = (threadgroup const float4*)p; + for (int j = 0; j < 4; ++j) { + const float4 a = v[j]; + xs[4 * j] = a.x; + xs[4 * j + 1] = a.y; + xs[4 * j + 2] = a.z; + xs[4 * j + 3] = a.w; + } +} +static inline void load16(device const float* p, thread float* xs) { + device const float4* v = (device const float4*)p; + for (int j = 0; j < 4; ++j) { + const float4 a = v[j]; + xs[4 * j] = a.x; + xs[4 * j + 1] = a.y; + xs[4 * j + 2] = a.z; + xs[4 * j + 3] = a.w; + } +} + +// ---- prepared x (xprep): XP holds, per token row of n floats, the kinds +// XK_RAW x (every format without the mask trick) +// XK_Q2 x * 4^-(i % 8) (Q2_0: 2-bit fields in place in a 16-bit word) +// XK_P4 x * 2^-(8 (i % 4)) (Q4_K, Q5_K: byte fields in place in a 32-bit word) +// XK_P2 x * 2^-(8 (i % 2)) (Q6_K, Q3_K: byte fields in place in a 16-bit word) +// and S16 = the sum of every 16 consecutive raw x (the formats with an offset or a min). +#define XK_RAW 0 +#define XK_Q2 1 +#define XK_P4 2 +#define XK_P2 3 +template struct XKind { enum { K = XK_RAW }; }; +template <> struct XKind { enum { K = XK_Q2 }; }; +template <> struct XKind { enum { K = XK_P4 }; }; +template <> struct XKind { enum { K = XK_P4 }; }; +template <> struct XKind { enum { K = XK_P2 }; }; +template <> struct XKind { enum { K = XK_P2 }; }; + +// X::load(xp, s16, u, xs, xsum): unit u's prepared x (xp = the token's row of the format's kind, s16 = its sums) +template struct X { + template static inline void load(P xp, P, int u, thread float* xs, thread float* xsum) { + constexpr int UNIT = fmt::UNIT; + for (int k = 0; k < UNIT; k += 16) load16(xp + u * UNIT + k, xs + k); + xsum[0] = xsum[1] = 0.0f; + } +}; +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + load16(xp + 16 * u, xs); + xsum[0] = s16[u]; + xsum[1] = 0.0f; + } +}; +// Q4_K / Q5_K: unit (block b, chunk c, half m): run0 at 256b + 64c + 16m, run1 = run0 + 32 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + const int p0 = 256 * (u >> 3) + 64 * ((u >> 1) & 3) + 16 * (u & 1); + load16(xp + p0, xs); + load16(xp + p0 + 32, xs + 16); + xsum[0] = s16[p0 >> 4]; + xsum[1] = s16[(p0 >> 4) + 2]; + } +}; +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + X::load(xp, s16, u, xs, xsum); + } +}; +// Q6_K: unit (block b, half n, m, p): run0 at 256b + 128n + 32p + 16m, run1 = run0 + 64 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + const int p0 = 256 * (u >> 3) + 128 * ((u >> 2) & 1) + 32 * (u & 1) + 16 * ((u >> 1) & 1); + load16(xp + p0, xs); + load16(xp + p0 + 64, xs + 16); + xsum[0] = s16[p0 >> 4]; + xsum[1] = s16[(p0 >> 4) + 4]; + } +}; +// Q5_1: 32 contiguous raw x, the two per-16 sums formed from them (the GEMVs that read Q5_1 - the routed down - get +// h without prepared sums) +template <> struct X { + template static inline void load(P xp, P, int u, thread float* xs, thread float* xsum) { + load16(xp + 32 * u, xs); + load16(xp + 32 * u + 16, xs + 16); + float a = 0.0f, b = 0.0f; + for (int j = 0; j < 16; ++j) { a += xs[j]; b += xs[16 + j]; } + xsum[0] = a; + xsum[1] = b; + } +}; +// Q3_K: 32 contiguous, sums per 16 +template <> struct X { + template static inline void load(P xp, P s16, int u, thread float* xs, thread float* xsum) { + load16(xp + 32 * u, xs); + load16(xp + 32 * u + 16, xs + 16); + xsum[0] = s16[2 * u]; + xsum[1] = s16[2 * u + 1]; + } +}; + +// xprep: one thread per 16 consecutive values of one token row; XP = [4 kinds][n_tok][n] then S16 [n_tok][n / 16] +kernel void xprep(device const float* x [[buffer(0)]], device float* xp [[buffer(1)]], constant uint& n [[buffer(2)]], + constant uint& n_tok [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const uint per = n / 16; + if (gid >= per * n_tok) return; + const uint t = gid / per, i = gid % per; + device const float* src = x + t * n + 16 * i; + const uint plane = n_tok * n; + device float* raw = xp + t * n + 16 * i; + float v[16]; + load16(src, v); + float s0 = 0, s1 = 0; + for (int j = 0; j < 16; j += 2) { s0 += v[j]; s1 += v[j + 1]; } + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; + const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + for (int j = 0; j < 16; ++j) { + raw[j] = v[j]; + raw[plane + j] = v[j] * q2s[j & 7]; + raw[2 * plane + j] = v[j] * p4s[j & 3]; + raw[3 * plane + j] = (j & 1) ? v[j] * 0.00390625f : v[j]; + } + xp[4 * plane + t * per + i] = s0 + s1; +} + +// ---- decoders: Dec::run(row, u, q, sc) decodes unit u of a row into q[UNIT] and sc = (d1, d2, m1, m2), so that +// for any token: dot = d1 * sum_{j < UNIT/2} q[j] xs[j] + d2 * sum_{j >= UNIT/2} q[j] xs[j] - (m1 xsum0 + m2 xsum1) +// with xs, xsum from X::load. A unit's weights are decoded once per row and applied to every token. +template struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint* tab) { + constexpr int UNIT = fmt::UNIT, UPB = fmt::QK / UNIT; + const float s = deq(row + (u / UPB) * fmt::BB, u % UPB, q, tab); + sc = float4(s, s, 0.0f, 0.0f); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float(a & ushort(3u << (2 * j))); + q[8 + j] = float(c & ushort(3u << (2 * j))); + } + const float d = rdh(b); + sc = float4(d, d, d, 0.0f); + } +}; + +// Q2_0 with raw x: q = c - 1 (shift, and, convert, subtract), one scale +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 2) * 18; + device const ushort* p = (device const ushort*)(b + 2 + 4 * (u & 3)); + const ushort a = p[0], c = p[1]; + for (int j = 0; j < 8; ++j) { + q[j] = float((a >> (2 * j)) & 3) - 1.0f; + q[8 + j] = float((c >> (2 * j)) & 3) - 1.0f; + } + const float d = rdh(b); + sc = float4(d, d, 0.0f, 0.0f); + } +}; + +static inline void scale_min_k4(int j, device const uchar* q, thread float& d, thread float& m) { + if (j < 4) { + d = float(q[j] & 63); + m = float(q[j + 4] & 63); + } else { + d = float((q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4)); + m = float((q[j + 4] >> 4) | ((q[j] >> 6) << 4)); + } +} +static inline void bytes4(uint w, thread float* q) { + q[0] = float(w & 0xFFu); + q[1] = float(w & 0xFF00u); + q[2] = float(w & 0xFF0000u); + q[3] = float(w & 0xFF000000u); +} +static inline void bytes2(ushort w, thread float* q) { + q[0] = float(w & 0xFF); + q[1] = float(w & 0xFF00); +} + +template <> struct Dec { // w = d q + m: sum = d (q . x) + m (sum x) + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint* tab) { + device const uchar* b = row + u * 24; + const float d = deq(b, 0, q, tab), m = rdh(b + 2); + sc = float4(d, d, -m, -m); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 144; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + device const uint* qp = (device const uint*)(b + 16 + 32 * c + 16 * m); + for (int w = 0; w < 4; ++w) { + const uint v = qp[w]; + bytes4(v & 0x0F0F0F0Fu, q + 4 * w); + bytes4((v >> 4) & 0x0F0F0F0Fu, q + 16 + 4 * w); + } + sc = float4(d * s1, d * s2, dm * m1, dm * m2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 176; + const int c = (u >> 1) & 3, m = u & 1; + const float d = rdh(b), dm = rdh(b + 2); + float s1, m1, s2, m2; + scale_min_k4(2 * c, b + 4, s1, m1); + scale_min_k4(2 * c + 1, b + 4, s2, m2); + device const uint* qp = (device const uint*)(b + 48 + 32 * c + 16 * m); + device const uint* hp = (device const uint*)(b + 16 + 16 * m); + for (int w = 0; w < 4; ++w) { + const uint v = qp[w], h = hp[w]; + bytes4((v & 0x0F0F0F0Fu) | (((h >> (2 * c)) & 0x01010101u) << 4), q + 4 * w); + bytes4(((v >> 4) & 0x0F0F0F0Fu) | (((h >> (2 * c + 1)) & 0x01010101u) << 4), q + 16 + 4 * w); + } + sc = float4(d * s1, d * s2, dm * m1, dm * m2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 210; + const int n = (u >> 2) & 1, m = (u >> 1) & 1, p = u & 1; + device const ushort* ql = (device const ushort*)(b + 64 * n + 32 * p + 16 * m); + device const ushort* qh = (device const ushort*)(b + 128 + 32 * n + 16 * m); + device const char* scs = (device const char*)(b + 192); + const float d = rdh(b + 208); + for (int w = 0; w < 8; ++w) { + const ushort l = ql[w], h = qh[w]; + bytes2((l & 0x0F0F) | (((h >> (2 * p)) & 0x0303) << 4), q + 2 * w); + bytes2(((l >> 4) & 0x0F0F) | (((h >> (2 * p + 4)) & 0x0303) << 4), q + 16 + 2 * w); + } + const float d1 = d * float(scs[8 * n + m + 2 * p]), d2 = d * float(scs[8 * n + m + 2 * p + 4]); + sc = float4(d1, d2, 32.0f * d1, 32.0f * d2); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uchar* b = row + (u >> 3) * 110; + const int n = (u >> 2) & 1, j = u & 3; + // qs and hmask as packed (2-byte aligned) vectors: blocks are 110 bytes, so odd blocks are only 2-aligned + device const packed_ushort4* qv = (device const packed_ushort4*)(b + 32 + 32 * n); + device const packed_ushort4* hv = (device const packed_ushort4*)b; + ushort qq[16], hm[16]; + for (int v = 0; v < 4; ++v) { + const ushort4 a4 = ushort4(qv[v]), h4 = ushort4(hv[v]); + qq[4 * v] = a4.x; qq[4 * v + 1] = a4.y; qq[4 * v + 2] = a4.z; qq[4 * v + 3] = a4.w; + hm[4 * v] = h4.x; hm[4 * v + 1] = h4.y; hm[4 * v + 2] = h4.z; hm[4 * v + 3] = h4.w; + } + device const uchar* s = b + 96; + const int i1 = 8 * n + 2 * j, i2 = i1 + 1; + const int l1 = (i1 < 8 ? s[i1] & 0xF : s[i1 - 8] >> 4) | (((s[8 + (i1 & 3)] >> (2 * (i1 >> 2))) & 3) << 4); + const int l2 = (i2 < 8 ? s[i2] & 0xF : s[i2 - 8] >> 4) | (((s[8 + (i2 & 3)] >> (2 * (i2 >> 2))) & 3) << 4); + const float d = rdh(b + 108); + const int hs = 4 * n + j; + for (int w = 0; w < 8; ++w) { + bytes2(((qq[w] >> (2 * j)) & 0x0303) | (((hm[w] >> hs) & 0x0101) << 2), q + 2 * w); + bytes2(((qq[w + 8] >> (2 * j)) & 0x0303) | (((hm[w + 8] >> hs) & 0x0101) << 2), q + 16 + 2 * w); + } + const float d1 = d * float(l1 - 32), d2 = d * float(l2 - 32); + sc = float4(d1, d2, 4.0f * d1, 4.0f * d2); + } +}; + +template <> struct Dec { // 16 floats a unit (the Unsloth files' router / alpha / beta / inject rows) + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const float4* p = (device const float4*)(row + 64 * u); + for (int k = 0; k < 4; ++k) { + const float4 v = p[k]; + q[4 * k] = v.x; q[4 * k + 1] = v.y; q[4 * k + 2] = v.z; q[4 * k + 3] = v.w; + } + sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + } +}; + +template <> struct Dec { + static inline void run(device const uchar* row, int u, thread float* q, thread float4& sc, threadgroup const uint*) { + device const uint4* p = (device const uint4*)(row + 32 * u); + const uint4 a = p[0], c = p[1]; + const uint v[8] = {a.x, a.y, a.z, a.w, c.x, c.y, c.z, c.w}; + for (int k = 0; k < 8; ++k) { + q[2 * k] = as_type(v[k] << 16); + q[2 * k + 1] = as_type(v[k] & 0xFFFF0000u); + } + sc = float4(1.0f, 1.0f, 0.0f, 0.0f); + } +}; + +// acc += unit dot of decoded q / sc with one token's prepared slice +template +static inline float unit_combine(thread const float* q, float4 sc, thread const float* xs, thread const float* xsum) { + constexpr int H = UNIT / 2; + float s0 = 0, s1 = 0, s2 = 0, s3 = 0; + for (int j = 0; j < H; j += 2) { + s0 = fma(q[j], xs[j], s0); + s1 = fma(q[j + 1], xs[j + 1], s1); + s2 = fma(q[H + j], xs[H + j], s2); + s3 = fma(q[H + j + 1], xs[H + j + 1], s3); + } + return sc.x * (s0 + s1) + sc.y * (s2 + s3) - (sc.z * xsum[0] + sc.w * xsum[1]); +} diff --git a/edge0-strata/metal/kernels/wrap_probe.metal b/edge0-strata/metal/kernels/wrap_probe.metal new file mode 100644 index 00000000..1cc2a7a1 --- /dev/null +++ b/edge0-strata/metal/kernels/wrap_probe.metal @@ -0,0 +1,12 @@ +// wrap_probe.metal - the tests/wrap_split_test.cpp kernel: reads a window at the start of a bound range and another at +// its far end (`far` floats on), to prove a chunked no-copy buffer (Context::wrap with a span) binds the whole range. +#include +using namespace metal; + +kernel void probe_copy(device const float* src [[buffer(0)]], device float* dst [[buffer(1)]], + device const uint* par [[buffer(2)]], uint i [[thread_position_in_grid]]) { + const uint n = par[0], far = par[1]; + if (i >= n) return; + dst[i] = src[i]; + dst[n + i] = src[far + i]; +} diff --git a/edge0-strata/metal/probe/mkfile.c b/edge0-strata/metal/probe/mkfile.c new file mode 100644 index 00000000..068aadbd --- /dev/null +++ b/edge0-strata/metal/probe/mkfile.c @@ -0,0 +1,23 @@ +// mkfile.c - write a dense test file of N GiB with F_NOCACHE (so it is not left in the page cache): +// the first uint32 of every 4 KiB is (offset >> 12) ^ 0x5a5a5a5a, the rest is a cheap pattern. +#include +#include +#include +#include +#include +int main(int argc, char** argv) { + if (argc < 3) return 1; + const uint64_t gib = strtoull(argv[2], 0, 10); + int fd = open(argv[1], O_WRONLY | O_CREAT | O_TRUNC, 0644); + if (fd < 0) return 2; + fcntl(fd, F_NOCACHE, 1); + const size_t B = 64u << 20; + uint32_t* buf = aligned_alloc(16384, B); + for (uint64_t off = 0; off < (gib << 30); off += B) { + for (size_t i = 0; i < B / 4; ++i) buf[i] = (uint32_t) (i * 2654435761u); + for (size_t p = 0; p < B; p += 4096) buf[p / 4] = (uint32_t) ((off + p) >> 12) ^ 0x5a5a5a5au; + if (write(fd, buf, B) != (ssize_t) B) return 3; + } + close(fd); + return 0; +} diff --git a/edge0-strata/metal/probe/mtl_doorbell.cpp b/edge0-strata/metal/probe/mtl_doorbell.cpp new file mode 100644 index 00000000..20baf0d8 --- /dev/null +++ b/edge0-strata/metal/probe/mtl_doorbell.cpp @@ -0,0 +1,115 @@ +// mtl_doorbell.cpp - can Strata's doorbell protocol run on Metal? +// A. in-kernel spin: a running kernel polls a shared-memory word that the CPU writes later (the CUDA doorbell model) +// B. MTLSharedEvent ping-pong: the GPU signals "routing ready" (value 2l+1), the host sees it, does its part, sets +// 2l+2 from the CPU, and the GPU (encodeWait inside the same, already committed command buffer) goes on. +// Measures the round trip per "layer". +#define NS_PRIVATE_IMPLEMENTATION +#define MTL_PRIVATE_IMPLEMENTATION +#include +#include + +#include +#include +#include +#include +#include +#include + +static double now_us() { + return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); +} +static const char* kSrc = R"MSL( +#include +using namespace metal; +struct SpinArgs { device atomic_uint* flag; device uint* out; uint want; uint max_iter; }; +kernel void spin(constant SpinArgs& p [[buffer(0)]], uint i [[thread_position_in_grid]]) { + if (i != 0) return; + uint it = 0; + while (atomic_load_explicit(p.flag, memory_order_relaxed) != p.want && it < p.max_iter) ++it; + p.out[0] = it; + p.out[1] = atomic_load_explicit(p.flag, memory_order_relaxed); +} +struct TickArgs { device uint* x; }; +kernel void tick(constant TickArgs& p [[buffer(0)]], uint i [[thread_position_in_grid]]) { if (i == 0) p.x[0] += 1; } +)MSL"; +struct SpinArgs { uint64_t flag, out; uint32_t want, max_iter; }; +struct TickArgs { uint64_t x; }; + +int main() { + NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init(); + MTL::Device* dev = MTL::CreateSystemDefaultDevice(); + NS::Error* e = nullptr; + MTL::Library* lib = dev->newLibrary(NS::String::string(kSrc, NS::UTF8StringEncoding), nullptr, &e); + if (!lib) { std::printf("compile failed\n"); return 1; } + auto pipe = [&](const char* n) { return dev->newComputePipelineState(lib->newFunction(NS::String::string(n, NS::UTF8StringEncoding)), &e); }; + MTL::ComputePipelineState* p_spin = pipe("spin"); + MTL::ComputePipelineState* p_tick = pipe("tick"); + MTL::CommandQueue* q = dev->newCommandQueue(); + MTL::Buffer* bf = dev->newBuffer(16384, MTL::ResourceStorageModeShared); + MTL::Buffer* bo = dev->newBuffer(16384, MTL::ResourceStorageModeShared); + MTL::ResidencySet* rs = dev->newResidencySet(MTL::ResidencySetDescriptor::alloc()->init(), &e); + rs->addAllocation(bf); rs->addAllocation(bo); rs->commit(); q->addResidencySet(rs); + volatile uint32_t* flag = (volatile uint32_t*) bf->contents(); + uint32_t* out = (uint32_t*) bo->contents(); + + // ---- A: in-kernel spin on a word the CPU writes 5 ms after commit + for (int rep = 0; rep < 3; ++rep) { + *flag = 0; out[0] = out[1] = 0; + SpinArgs a{bf->gpuAddress(), bo->gpuAddress(), 7, 200000000u}; + MTL::CommandBuffer* cb = q->commandBuffer(); + MTL::ComputeCommandEncoder* enc = cb->computeCommandEncoder(); + enc->setComputePipelineState(p_spin); + enc->setBytes(&a, sizeof a, 0); + enc->dispatchThreads(MTL::Size(1, 1, 1), MTL::Size(1, 1, 1)); + enc->endEncoding(); + const double t0 = now_us(); + cb->commit(); + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + std::atomic_thread_fence(std::memory_order_seq_cst); + *flag = 7; + const double tw = now_us(); + cb->waitUntilCompleted(); + std::printf("A in-kernel spin: GPU saw the CPU store: %s (iterations %u, flag %u), done %.0f us after the store, status %d\n", + out[1] == 7 ? "YES" : "NO (timed out)", out[0], out[1], now_us() - tw, (int) cb->status()); + (void) t0; + } + + // ---- B: shared-event ping-pong, L "layers" in ONE committed command buffer + for (int rep = 0; rep < 3; ++rep) { + const int L = 48; + MTL::SharedEvent* ev = dev->newSharedEvent(); + ev->setSignaledValue(0); + out[0] = 0; + MTL::CommandBuffer* cb = q->commandBuffer(); + TickArgs ta{bo->gpuAddress()}; + for (int l = 0; l < L; ++l) { + MTL::ComputeCommandEncoder* enc = cb->computeCommandEncoder(); + enc->setComputePipelineState(p_tick); + enc->setBytes(&ta, sizeof ta, 0); + enc->dispatchThreads(MTL::Size(1, 1, 1), MTL::Size(1, 1, 1)); + enc->endEncoding(); + cb->encodeSignalEvent(ev, 2 * l + 1); // "the miss list of layer l is ready" + cb->encodeWait(ev, 2 * l + 2); // "the host has answered layer l" + } + std::vector rt; + const double t0 = now_us(); + cb->commit(); + for (int l = 0; l < L; ++l) { + while (ev->signaledValue() < (uint64_t) (2 * l + 1)) {} + const double ts = now_us(); + if (out[0] != (uint32_t) (l + 1)) std::printf(" layer %d: host read %u (stale?)\n", l, out[0]); + ev->setSignaledValue(2 * l + 2); + rt.push_back(ts); + } + cb->waitUntilCompleted(); + const double total = now_us() - t0; + std::vector gaps; + for (int l = 1; l < L; ++l) gaps.push_back(rt[l] - rt[l - 1]); + std::sort(gaps.begin(), gaps.end()); + std::printf("B shared-event round trip: %d layers in %.0f us, per layer median %.1f us (min %.1f, max %.1f), ticks %u, status %d\n", + L, total, gaps[gaps.size() / 2], gaps.front(), gaps.back(), out[0], (int) cb->status()); + ev->release(); + } + pool->release(); + return 0; +} diff --git a/edge0-strata/metal/probe/mtl_probe.cpp b/edge0-strata/metal/probe/mtl_probe.cpp new file mode 100644 index 00000000..d1e0a335 --- /dev/null +++ b/edge0-strata/metal/probe/mtl_probe.cpp @@ -0,0 +1,268 @@ +// mtl_probe.cpp - M5 Max checks for the Metal device-interface design (mac-design stream). +// 1. device facts: unified memory, maxBufferLength, recommendedMaxWorkingSetSize, GPU families +// 2. MSL compiled at run time (no offline `metal` tool in the Command Line Tools) +// 3. a hello kernel that takes RAW GPU ADDRESSES in a setBytes struct (the CUDA "void* device pointer" model) +// 4. gpuAddress vs contents(): are they equal? +// 5. newBufferWithBytesNoCopy over an mmapped file > 30 GB: alignment, max length, split into chunks, GPU reads it +// Build: clang++ -std=c++20 -O2 -I metal-cpp mtl_probe.cpp -framework Metal -framework Foundation -framework QuartzCore +#define NS_PRIVATE_IMPLEMENTATION +#define MTL_PRIVATE_IMPLEMENTATION +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +static double now_ms() { + return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +static const char* kSrc = R"MSL( +#include +using namespace metal; +struct AxpyArgs { device const float* x; device float* y; float a; uint n; }; +kernel void axpy(constant AxpyArgs& p [[buffer(0)]], uint i [[thread_position_in_grid]]) { + if (i < p.n) p.y[i] = p.a * p.x[i] + p.y[i]; +} +// reads one uint32 at each probe offset of a (possibly huge) buffer given by raw GPU address +struct ProbeArgs { device const uint* base; device const ulong* offs; device uint* out; uint n; }; +kernel void probe(constant ProbeArgs& p [[buffer(0)]], uint i [[thread_position_in_grid]]) { + if (i < p.n) p.out[i] = p.base[p.offs[i] / 4]; +} +// streaming read of a buffer range: sum of uint32 (bandwidth through a NoCopy file buffer) +struct SumArgs { device const uint4* base; device atomic_uint* out; ulong n4; }; +kernel void sum(constant SumArgs& p [[buffer(0)]], uint i [[thread_position_in_grid]], uint g [[threads_per_grid]]) { + uint s = 0; + for (ulong k = i; k < p.n4; k += g) { uint4 v = p.base[k]; s += v.x ^ v.y ^ v.z ^ v.w; } + atomic_fetch_add_explicit(p.out, s, memory_order_relaxed); +} +)MSL"; + +struct AxpyArgs { uint64_t x, y; float a; uint32_t n; }; +struct ProbeArgs { uint64_t base, offs, out; uint32_t n; }; +struct SumArgs { uint64_t base, out; uint64_t n4; }; + +int main(int argc, char** argv) { + NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init(); + MTL::Device* dev = MTL::CreateSystemDefaultDevice(); + if (!dev) { std::printf("no Metal device\n"); return 1; } + std::printf("device: %s\n", dev->name()->utf8String()); + std::printf("unified memory: %d\n", (int) dev->hasUnifiedMemory()); + std::printf("maxBufferLength: %llu (%.2f GiB)\n", (unsigned long long) dev->maxBufferLength(), dev->maxBufferLength() / 1073741824.0); + std::printf("recommendedMaxWorkingSetSize: %llu (%.2f GiB)\n", (unsigned long long) dev->recommendedMaxWorkingSetSize(), + dev->recommendedMaxWorkingSetSize() / 1073741824.0); + std::printf("maxThreadgroupMemoryLength: %llu\n", (unsigned long long) dev->maxThreadgroupMemoryLength()); + std::printf("families: apple7 %d apple8 %d apple9 %d apple10 %d metal3 %d metal4 %d\n", (int) dev->supportsFamily(MTL::GPUFamilyApple7), + (int) dev->supportsFamily(MTL::GPUFamilyApple8), (int) dev->supportsFamily(MTL::GPUFamilyApple9), + (int) dev->supportsFamily(MTL::GPUFamilyApple10), (int) dev->supportsFamily(MTL::GPUFamilyMetal3), + (int) dev->supportsFamily(MTL::GPUFamilyMetal4)); + std::printf("page size: %ld\n", sysconf(_SC_PAGESIZE)); + + // ---- 2. run-time MSL compile + NS::Error* e = nullptr; + double t0 = now_ms(); + MTL::CompileOptions* co = MTL::CompileOptions::alloc()->init(); + MTL::Library* lib = dev->newLibrary(NS::String::string(kSrc, NS::UTF8StringEncoding), co, &e); + if (!lib) { std::printf("MSL compile failed: %s\n", e ? e->localizedDescription()->utf8String() : "?"); return 1; } + std::printf("MSL run-time compile: %.1f ms\n", now_ms() - t0); + auto pipe = [&](const char* n) { + MTL::Function* f = lib->newFunction(NS::String::string(n, NS::UTF8StringEncoding)); + MTL::ComputePipelineState* p = dev->newComputePipelineState(f, &e); + if (!p) { std::printf("pipeline %s failed\n", n); std::exit(1); } + return p; + }; + MTL::ComputePipelineState* p_axpy = pipe("axpy"); + MTL::ComputePipelineState* p_probe = pipe("probe"); + MTL::ComputePipelineState* p_sum = pipe("sum"); + MTL::CommandQueue* q = dev->newCommandQueue(); + + // ---- 3/4. hello kernel through raw GPU addresses; residency through a residency set + const uint32_t N = 1 << 20; + MTL::Buffer* bx = dev->newBuffer(N * 4, MTL::ResourceStorageModeShared); + MTL::Buffer* by = dev->newBuffer(N * 4, MTL::ResourceStorageModeShared); + std::printf("buffer x: contents %p gpuAddress 0x%llx (%s)\n", bx->contents(), (unsigned long long) bx->gpuAddress(), + (uint64_t) bx->contents() == bx->gpuAddress() ? "EQUAL" : "different"); + float* hx = (float*) bx->contents(); + float* hy = (float*) by->contents(); + for (uint32_t i = 0; i < N; ++i) { hx[i] = (float) i; hy[i] = 1.0f; } + MTL::ResidencySetDescriptor* rd = MTL::ResidencySetDescriptor::alloc()->init(); + MTL::ResidencySet* rs = dev->newResidencySet(rd, &e); + if (!rs) { std::printf("residency set failed\n"); return 1; } + rs->addAllocation(bx); rs->addAllocation(by); rs->commit(); + q->addResidencySet(rs); + { + AxpyArgs a{bx->gpuAddress(), by->gpuAddress(), 2.0f, N}; + MTL::CommandBuffer* cb = q->commandBuffer(); + MTL::ComputeCommandEncoder* enc = cb->computeCommandEncoder(); + enc->setComputePipelineState(p_axpy); + enc->setBytes(&a, sizeof a, 0); + enc->dispatchThreads(MTL::Size(N, 1, 1), MTL::Size(256, 1, 1)); + enc->endEncoding(); + cb->commit(); + cb->waitUntilCompleted(); + int bad = 0; + for (uint32_t i = 0; i < N; ++i) if (hy[i] != 2.0f * i + 1.0f) ++bad; + std::printf("hello axpy via raw GPU addresses + residency set: %s (%d bad)\n", bad ? "FAIL" : "OK", bad); + } + // MTLSharedEvent as the CUDA event: signal from the GPU, poll from the host + { + MTL::SharedEvent* ev = dev->newSharedEvent(); + MTL::CommandBuffer* cb = q->commandBuffer(); + MTL::ComputeCommandEncoder* enc = cb->computeCommandEncoder(); + AxpyArgs a{bx->gpuAddress(), by->gpuAddress(), 1.0f, N}; + enc->setComputePipelineState(p_axpy); + enc->setBytes(&a, sizeof a, 0); + enc->dispatchThreads(MTL::Size(N, 1, 1), MTL::Size(256, 1, 1)); + enc->endEncoding(); + cb->encodeSignalEvent(ev, 1); + double ts = now_ms(); + cb->commit(); + long spins = 0; + while (ev->signaledValue() < 1) ++spins; + std::printf("shared event: signalled after %.3f ms, %ld host polls\n", now_ms() - ts, spins); + ev->release(); + } + // dispatch cost: 1,000 tiny dispatches in one command buffer (the window-replay cost model) + { + AxpyArgs a{bx->gpuAddress(), by->gpuAddress(), 0.0f, 256}; + for (int rep = 0; rep < 3; ++rep) { + double ts = now_ms(); + MTL::CommandBuffer* cb = q->commandBufferWithUnretainedReferences(); + MTL::ComputeCommandEncoder* enc = cb->computeCommandEncoder(); + for (int k = 0; k < 1000; ++k) { + enc->setComputePipelineState(p_axpy); + enc->setBytes(&a, sizeof a, 0); + enc->dispatchThreads(MTL::Size(256, 1, 1), MTL::Size(256, 1, 1)); + } + enc->endEncoding(); + const double te = now_ms(); + cb->commit(); + cb->waitUntilCompleted(); + std::printf("1000 dependent tiny dispatches: encode %.3f ms, encode+run %.3f ms\n", te - ts, now_ms() - ts); + } + } + + // ---- 5. a > 30 GB mmapped file wrapped with newBufferWithBytesNoCopy + if (argc < 2) { pool->release(); return 0; } + const char* path = argv[1]; + const int fd = open(path, O_RDONLY); + if (fd < 0) { std::printf("open %s failed\n", path); return 1; } + struct stat sb; fstat(fd, &sb); + const uint64_t fsize = (uint64_t) sb.st_size; + const uint64_t pg = (uint64_t) sysconf(_SC_PAGESIZE); + const uint64_t mlen = (fsize + pg - 1) / pg * pg; // NoCopy wants a page-multiple length + std::printf("file %s: %.2f GiB\n", path, fsize / 1073741824.0); + for (int mode = 0; mode < 2; ++mode) { + // mode 0: PROT_READ MAP_SHARED (read-only, what the engine wants); mode 1: PROT_READ|PROT_WRITE MAP_PRIVATE + const int prot = mode == 0 ? PROT_READ : (PROT_READ | PROT_WRITE); + const int flags = mode == 0 ? MAP_SHARED : MAP_PRIVATE; + void* m = mmap(nullptr, mlen, prot, flags, fd, 0); + if (m == MAP_FAILED) { std::printf("mmap mode %d failed\n", mode); continue; } + std::printf("-- mode %d (%s): mmap at %p (page aligned %d)\n", mode, mode == 0 ? "PROT_READ MAP_SHARED" : "RW MAP_PRIVATE", m, + (int) (((uintptr_t) m % pg) == 0)); + double ts = now_ms(); + MTL::Buffer* whole = dev->newBuffer(m, mlen, MTL::ResourceStorageModeShared, nullptr); + std::printf(" one NoCopy buffer of %.2f GiB: %s (%.1f ms)\n", mlen / 1073741824.0, whole ? "OK" : "nil", now_ms() - ts); + // chunks of at most maxBufferLength / 2 (and a 4 GiB variant: one chunk per few layers) + std::vector chunks; + const uint64_t chunk = argc > 2 ? (std::strtoull(argv[2], nullptr, 10) << 20) : (4ull << 30); + ts = now_ms(); + bool all = true; + for (uint64_t off = 0; off < mlen; off += chunk) { + const uint64_t len = std::min(chunk, mlen - off); + MTL::Buffer* b = dev->newBuffer((char*) m + off, len, MTL::ResourceStorageModeShared, nullptr); + if (!b) { all = false; break; } + chunks.push_back(b); + } + std::printf(" %zu NoCopy chunks of %.2f GiB: %s (%.1f ms)\n", chunks.size(), chunk / 1073741824.0, all ? "OK" : "FAILED", now_ms() - ts); + // GPU probe reads at 64 offsets spread over the file (the file's pattern: uint32 at offset o equals (o>>12) ^ 0x5a5a5a5a + // where written; holes read 0) + const uint32_t NP = 64; + MTL::Buffer* boffs = dev->newBuffer(NP * 8, MTL::ResourceStorageModeShared); + MTL::Buffer* bout = dev->newBuffer(NP * 4, MTL::ResourceStorageModeShared); + uint64_t* offs = (uint64_t*) boffs->contents(); + for (uint32_t i = 0; i < NP; ++i) offs[i] = ((fsize / NP) * i) & ~(uint64_t) 4095; + auto run_probe = [&](MTL::Buffer* b, uint64_t base_off, const char* what) { + // offsets relative to this buffer's base + std::vector keep(offs, offs + NP); + uint32_t n = 0; + std::vector idx; + for (uint32_t i = 0; i < NP; ++i) + if (keep[i] >= base_off && keep[i] < base_off + b->length()) { offs[n++] = keep[i] - base_off; idx.push_back(i); } + MTL::ResidencySet* r2 = dev->newResidencySet(rd, &e); + r2->addAllocation(b); r2->addAllocation(boffs); r2->addAllocation(bout); r2->commit(); + double tr = now_ms(); + r2->requestResidency(); + const double treq = now_ms() - tr; + ProbeArgs a{b->gpuAddress(), boffs->gpuAddress(), bout->gpuAddress(), n}; + MTL::CommandBuffer* cb = q->commandBuffer(); + cb->useResidencySet(r2); + MTL::ComputeCommandEncoder* enc = cb->computeCommandEncoder(); + enc->setComputePipelineState(p_probe); + enc->setBytes(&a, sizeof a, 0); + enc->dispatchThreads(MTL::Size(n ? n : 1, 1, 1), MTL::Size(64, 1, 1)); + enc->endEncoding(); + tr = now_ms(); + cb->commit(); + cb->waitUntilCompleted(); + const double trun = now_ms() - tr; + int okc = 0, bad = 0; + const uint32_t* out = (const uint32_t*) bout->contents(); + for (uint32_t k = 0; k < n; ++k) { + const uint64_t fo = keep[idx[k]]; + uint32_t want = 0; + pread(fd, &want, 4, (off_t) fo); + (out[k] == want ? okc : bad)++; + } + std::printf(" GPU probe %s: %u reads, %d match pread, %d differ (requestResidency %.1f ms, run %.1f ms, status %d)\n", what, n, + okc, bad, treq, trun, (int) cb->status()); + for (uint32_t i = 0; i < NP; ++i) offs[i] = keep[i]; + r2->release(); + }; + if (whole) run_probe(whole, 0, "whole buffer"); + if (all && !chunks.empty()) run_probe(chunks.back(), (chunks.size() - 1) * chunk, "last chunk"); + // a streaming GPU read over one chunk, cold then warm (page cache) + if (all && !chunks.empty()) { + MTL::Buffer* b = chunks[chunks.size() / 2]; + MTL::Buffer* bs = dev->newBuffer(4, MTL::ResourceStorageModeShared); + for (int rep = 0; rep < 2; ++rep) { + *(uint32_t*) bs->contents() = 0; + MTL::ResidencySet* r2 = dev->newResidencySet(rd, &e); + r2->addAllocation(b); r2->addAllocation(bs); r2->commit(); + SumArgs a{b->gpuAddress(), bs->gpuAddress(), b->length() / 16}; + double tr = now_ms(); + MTL::CommandBuffer* cb = q->commandBuffer(); + cb->useResidencySet(r2); + MTL::ComputeCommandEncoder* enc = cb->computeCommandEncoder(); + enc->setComputePipelineState(p_sum); + enc->setBytes(&a, sizeof a, 0); + enc->dispatchThreads(MTL::Size(40 * 1024, 1, 1), MTL::Size(256, 1, 1)); + enc->endEncoding(); + cb->commit(); + cb->waitUntilCompleted(); + const double dt = now_ms() - tr; + std::printf(" GPU streaming read of a %.2f GiB chunk (%s): %.1f ms = %.1f GB/s\n", b->length() / 1073741824.0, + rep == 0 ? "first touch" : "again", dt, b->length() / dt / 1e6); + r2->release(); + } + bs->release(); + } + for (auto* b : chunks) b->release(); + if (whole) whole->release(); + boffs->release(); bout->release(); + munmap(m, mlen); + } + close(fd); + pool->release(); + return 0; +} diff --git a/edge0-strata/metal/setup_mac.py b/edge0-strata/metal/setup_mac.py new file mode 100644 index 00000000..774ac3b2 --- /dev/null +++ b/edge0-strata/metal/setup_mac.py @@ -0,0 +1,578 @@ +"""setup.py on a Mac: Apple Silicon, the Metal engine (metal/), built from source on the Mac. + + ./setup.sh [options] (setup.py hands a macOS run to this file; the PC paths are not used) + +What it does, each step skipped when it is already done: + + 1. checks the Mac: Apple Silicon, the chip, unified memory, macOS version, the Command Line Tools, free disk + 2. recommends a model by memory (recommend, never force: an explicit --model is kept, with what it risks) + 3. installs the Python packages into .venv (cmake and ninja come from pip, just before step 4: no Homebrew, no sudo) + 4. builds the Metal engine from source: metal-cpp (Apple's headers) and llama.cpp's ggml at the pinned commit + 5. downloads the model (setup.py's downloader: resumable, the pinned revision) + 6. prepares the model: the native pack (tools/iq_pack.py) and the MTP draft layer (tools/mtp_fetch.py, mtp_pack.py) + 7. writes strata-.json and run-.sh with the Metal defaults (--spec 2, the Q8_0 hyper-connection + read) and starts the server (serve/server.py), the same as on a PC + +Everything shared - the model table, the downloader, the config and run-script writers, the server - is setup.py's. +""" +from __future__ import annotations + +import argparse +import atexit +import hashlib +import json +import os +import platform +import shutil +import subprocess +import sys +import zipfile +from pathlib import Path + +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) +import setup as S # noqa: E402 +from strata_control import mac_fit as _fit # noqa: E402 +from strata_control.mac_fit import ( # noqa: E402,F401 (re-exported: the tests and callers use these names here) + BIG_MAC_GB, CONTEXTS, GIB_GB, GPU_SHARE, MAC_ABOUT, MAC_MODELS, MIN_MACOS, OVERHEAD_GB, RECOMMENDED, + RESIDENT_LEAVE_SHARE, STREAM_COPY_GAIN, STREAM_MEASURED, STREAM_MTP_FROM_GB, STREAM_SPEC, TESTED_CONTEXT, + _sysctl, gpu_budget_gb, mac_info, mac_problems, macos_tuple, model_need_gb, model_note, recommend_context, + recommend_model, stream_mtp, streams) + +# Apple's metal-cpp headers for the macOS 26 SDK (https://developer.apple.com/metal/cpp/), pinned by SHA-256 +METAL_CPP_URL = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip" +METAL_CPP_SHA256 = "4df3c078b9aadcb516212e9cb03004cbc5ce9a3e9c068fa3144d021db585a3a4" +BUILD_DIR = ROOT / "build-metal" +EXE = BUILD_DIR / "strata" + +# the Unsloth files at setup.py's pinned revision: name -> (bytes, SHA-256), the Hub's LFS pointers (UD-Q4_K_XL's +# are setup.py's UNSLOTH_SHARDS) +UD_SHARDS = { + "Qwen3.8-Flash-Next-UD-Q5_K_XL-00001-of-00006.gguf": (10946618, "610922cdb1afe1094e95d7a3e86b12ee33317ebbf7320ac9344e21b6ae4bb69f"), + "Qwen3.8-Flash-Next-UD-Q5_K_XL-00002-of-00006.gguf": (682434912, "494ca4ed3dbf97bc28da88af3890b8877b9032f909812d00c0526a9ca5e91d2e"), + "Qwen3.8-Flash-Next-UD-Q5_K_XL-00003-of-00006.gguf": (54400261312, "34efd79a80a1ce540a517a5d56171924b66ce1c38b04c904f17ad6d8ef17cf20"), + "Qwen3.8-Flash-Next-UD-Q5_K_XL-00004-of-00006.gguf": (49990245824, "9ee6bbe462e6830864382d629101490625b1c95afa15492b7ad550c7bed2dd13"), + "Qwen3.8-Flash-Next-UD-Q5_K_XL-00005-of-00006.gguf": (49882416192, "0b2437f661f584030cb5aeb8fb1bb70f6422518a6746bfcc0304167c7b6b8beb"), + "Qwen3.8-Flash-Next-UD-Q5_K_XL-00006-of-00006.gguf": (3320101792, "55dd96c5ad0863b06b6393151b9ad0ac727b8d58a62e02a296ca5d5ddbfee4a3"), + "Qwen3.8-Flash-Next-UD-Q6_K_XL-00001-of-00006.gguf": (10946624, "8bdc6bad55b6699f46f60708829932ff4018d7ac9bd309c388534fa31753bdb7"), + "Qwen3.8-Flash-Next-UD-Q6_K_XL-00002-of-00006.gguf": (682434912, "494ca4ed3dbf97bc28da88af3890b8877b9032f909812d00c0526a9ca5e91d2e"), + "Qwen3.8-Flash-Next-UD-Q6_K_XL-00003-of-00006.gguf": (54400261312, "34efd79a80a1ce540a517a5d56171924b66ce1c38b04c904f17ad6d8ef17cf20"), + "Qwen3.8-Flash-Next-UD-Q6_K_XL-00004-of-00006.gguf": (49814108512, "0ea4b599880a5a52fcd9c188ba1443f646c69c96571904ddbd6a79246a01a997"), + "Qwen3.8-Flash-Next-UD-Q6_K_XL-00005-of-00006.gguf": (49419317344, "9948e81ae8368144b135a7403d8c0f2cf9a6c2f9f692313865855b2832e696c9"), + "Qwen3.8-Flash-Next-UD-Q6_K_XL-00006-of-00006.gguf": (14838313984, "e4ae255234f42b18012f6d6eec5bf615b57479b7e0a0233c47436a321ffeec3f"), +} + + +def model_shard_names(model: str) -> list[str]: + """The model's files in order (shard 1 first).""" + n = MAC_MODELS[model].get("shards", 2) + if MAC_MODELS[model].get("family") == "gyro": # agentionai's one file (the n-gram table is inside it) + return [S.model_file(S.FAMILIES["gyro"], model, 1)] + if MAC_MODELS[model].get("family") == "unsloth": + return [f"Qwen3.8-Flash-Next-{model}-{i:05d}-of-{n:05d}.gguf" for i in range(1, n + 1)] + return [S.model_file(S.FAMILIES["qwen"], model, i) for i in range(1, n + 1)] + + +def shard_sha(name: str): + return UD_SHARDS.get(name) or getattr(S, "UNSLOTH_SHARDS", {}).get(name) + + +# ------------------------------------------------------------------------------------------------ the Mac +# ------------------------------------------------------------------------------------------------ the engine +def metal_cpp() -> Path: + """Apple's metal-cpp headers in third_party/metal-cpp (one zip, pinned).""" + dst = ROOT / "third_party" / "metal-cpp" + if (dst / "Metal" / "Metal.hpp").exists(): + return dst + z = ROOT / "third_party" / "metal-cpp_26.zip" + S.download(METAL_CPP_URL, z, "metal-cpp (Apple's Metal headers)") + S.verify_sha256(z, z.stat().st_size, METAL_CPP_SHA256) + tmp = ROOT / "third_party" / "_metal_cpp" + shutil.rmtree(tmp, ignore_errors=True) + with zipfile.ZipFile(z) as f: + f.extractall(tmp) + shutil.rmtree(dst, ignore_errors=True) + shutil.move(str(tmp / "metal-cpp"), str(dst)) + shutil.rmtree(tmp, ignore_errors=True) + return dst + + +def tool(name: str) -> str: + """cmake / ninja from the venv's pip packages (no Homebrew).""" + p = Path(sys.executable).parent / name + return str(p) if p.exists() else (shutil.which(name) or name) + + +def engine_current() -> bool: + try: + meta = json.loads((BUILD_DIR / "BUILD.json").read_text(encoding="utf-8")) + except (OSError, ValueError): + return False + return EXE.exists() and meta.get("version") == S.source_version() and meta.get("source") == source_digest() + + +def source_digest() -> str: + """What the engine is built from: every source file's path, size and modification time (src/, include/, metal/, + third_party/ggml). A git pull, an unzipped update or a local edit changes it; nothing else does.""" + h = hashlib.sha256() + for d in ("src", "include", "metal", "third_party/ggml"): + for f in sorted((ROOT / d).rglob("*")): + if f.is_file() and f.suffix in (".cpp", ".cu", ".h", ".hpp", ".mm", ".metal", ".inc", ".txt", ".cuh"): + st = f.stat() + h.update(f"{f.relative_to(ROOT).as_posix()}|{st.st_size}|{int(st.st_mtime)}\n".encode()) + return h.hexdigest() + + +def build_engine(llama: Path, jobs: int | None = None) -> Path: + """The Metal engine from source: cmake + ninja (pip's), Apple's clang from the Command Line Tools.""" + if engine_current(): + S.ok(f"Metal engine: {EXE} (built for this source)") + return EXE + mcpp = metal_cpp() + S.ensure_build_tools() # cmake and ninja are in requirements-build.txt, not requirements.txt (#3403) + S.say(" Building the Metal engine from source (5-15 minutes the first time) ...") + BUILD_DIR.mkdir(exist_ok=True) + log = BUILD_DIR / "build.log" + env = dict(os.environ, PATH=str(Path(sys.executable).parent) + os.pathsep + os.environ.get("PATH", "")) + cfg = [tool("cmake"), "-S", str(ROOT / "metal"), "-B", str(BUILD_DIR), "-G", "Ninja", "-DCMAKE_BUILD_TYPE=Release", + f"-DCMAKE_MAKE_PROGRAM={tool('ninja')}", f"-DMETAL_CPP_DIR={mcpp}", f"-DSTRATA_GGML_DIR={llama}"] + build = [tool("ninja"), "-C", str(BUILD_DIR), "strata"] + (["-j", str(jobs)] if jobs else []) + with open(log, "w", encoding="utf-8") as f: + for cmd in (cfg, build): + rc = subprocess.call(cmd, stdout=f, stderr=subprocess.STDOUT, env=env) + if rc != 0: + S.fail(f"the Metal engine did not build (see {log})", + "the last lines of the log name the problem; the Command Line Tools must be up to date " + "(softwareupdate --list)") + (BUILD_DIR / "BUILD.json").write_text(json.dumps({"version": S.source_version(), "source": source_digest(), + "backend": "metal"}, indent=1), encoding="utf-8") + S.ok(f"Metal engine built: {EXE}") + return EXE + + +# ------------------------------------------------------------------------------------------------ the config +def engine_args(pack: Path, shards: list, rt: Path, ctx: int, stream: bool = False, ram_gb: float = 0.0) -> list[str]: + """The Metal defaults (strata_control/mac_fit.py), with this checkout's expert profile.""" + return _fit.engine_args(pack, shards, rt, ctx, stream, ram_gb, ROOT / "data" / "expert-profile.bin") + + +# the image encoder's projector (setup.py's pinned revision): name -> (bytes, SHA-256) +VISION_SHA256 = {"mmproj-Qwen3.8-Flash-Next-BF16.gguf": + (907543008, "b1a82259702816a5330d7bd7607cd9676b11780e79ff7348c21103ff3ce49bd0")} +VISION_EXE = ROOT / "build-vision" / "bin" / "strata-vision" + + +def build_vision(llama: Path) -> Path: + """strata-vision (tools/vision: llama.cpp's mtmd at the pinned commit), with ggml's Metal backend: the encoder + runs on the GPU.""" + stamp = ROOT / "build-vision" / "BUILD.json" + src = source_digest() + try: + if VISION_EXE.exists() and json.loads(stamp.read_text(encoding="utf-8")).get("source") == src: + S.ok(f"image encoder: {VISION_EXE}") + return VISION_EXE + except (OSError, ValueError): + pass + S.ensure_build_tools() + S.say(" Building the image encoder (1-3 minutes) ...") + env = dict(os.environ, PATH=str(Path(sys.executable).parent) + os.pathsep + os.environ.get("PATH", "")) + log = ROOT / "build-vision.log" + with open(log, "w", encoding="utf-8") as f: + for cmd in ([tool("cmake"), "-S", str(ROOT / "tools" / "vision"), "-B", str(ROOT / "build-vision"), "-G", "Ninja", + "-DCMAKE_BUILD_TYPE=Release", f"-DCMAKE_MAKE_PROGRAM={tool('ninja')}", f"-DLLAMA_DIR={llama}", + "-DSTRATA_VISION_CUDA=OFF", "-DSTRATA_VISION_METAL=ON", "-DSTRATA_PORTABLE=OFF"], + [tool("ninja"), "-C", str(ROOT / "build-vision"), "strata-vision"]): + if subprocess.call(cmd, stdout=f, stderr=subprocess.STDOUT, env=env) != 0: + S.fail(f"the image encoder did not build (see {log})") + stamp.write_text(json.dumps({"source": src}), encoding="utf-8") + S.ok(f"image encoder built: {VISION_EXE}") + return VISION_EXE + + +def ple_shard(files: list) -> Path: + """The shard that holds the n-gram table (per_layer_token_embd).""" + sys.path.insert(0, str(ROOT / "metal" / "tools")) + import metal_repack as R + for f in files: + rd = R.Reader(f) + rd.f.close() + if any(t["name"] == R.PLE for t in rd.tensors): + return f + S.fail(f"no shard of {files[0].name} holds the n-gram table (per_layer_token_embd)") + + +# the experts' streaming copy (metal/tools/stream_experts.py): what stays free on the disk after it is written +EXPERTS_MARGIN_GB = 10.0 + + +def experts_file_fits(need_gb: float, free_gb: float) -> bool: + return free_gb >= need_gb + EXPERTS_MARGIN_GB + + +def experts_file(native: Path) -> None: + """A streaming Mac: .experts beside the GPU's GGUF - every expert's rows in one piece, read with one SSD + request per expert instead of three (the engine uses it when it is there and matches; else the GGUF). Written + only when the disk has room for it and EXPERTS_MARGIN_GB more; otherwise setup says so and goes on.""" + out = native.with_name(native.name + ".experts") + if out.exists() and S.done(out): + S.ok(f"the experts' streaming copy: {out.name}") + return + sys.path.insert(0, str(ROOT / "metal" / "tools")) + import stream_experts as X + need = X.size_bytes(native) / 1e9 + free = S.free_gb(native.parent) + if not os.access(native.parent, os.W_OK): + S.warn(f"{native.parent} cannot be written: the experts' streaming copy ({need:.0f} GB) is not written; the " + f"engine streams from the model file ({STREAM_COPY_GAIN} slower)") + return + if not experts_file_fits(need, free): + S.warn(f"the experts' streaming copy needs {need:.0f} GB and {EXPERTS_MARGIN_GB:.0f} GB to spare on the disk " + f"({free:.0f} GB free): not written; the engine streams from the model file ({STREAM_COPY_GAIN} " + f"slower). Free some space and run ./setup.sh --setup again to add it") + return + S.say(f" Writing the experts in SSD-read order (one time, {need:.0f} GB; each expert is read with one request) ...") + S.run([sys.executable, str(ROOT / "metal" / "tools" / "stream_experts.py"), "--gguf", str(native), "--verify"]) + S.mark(out) + S.ok(f"the experts' streaming copy: {out.name}") + + +def mac_config(exe: Path, args: list, pack: Path, model_name: str, tag: str, port: int) -> dict: + return {"exe": str(exe), "args": args, "cwd": str(ROOT), "tokenizer": str(pack / "tokenizer"), + "model_name": model_name, "log": str(ROOT / f"strata-{tag.lower()}.log"), "lib_dirs": [], "port": port, + "backend": "metal"} + + +def mac_configs() -> list[Path]: + out = [] + for p in S.installed_configs(): + try: + if json.loads(p.read_text(encoding="utf-8-sig")).get("backend") == "metal": + out.append(p) + except (OSError, ValueError): + pass + return out + + +REPACK_LOCK = ".repack.lock" # metal_repack.py: .repack.lock holds the repack's pid while it runs +PREPARE_LOCK = ".preparing.lock" # setup: beside the repacked file, holds setup's pid from the repack to the end of the pack + + +def _pid_alive(pid: int) -> bool: + try: + os.kill(pid, 0) + except PermissionError: + return True # it exists, it is somebody else's + except OSError: + return False + return True + + +def repack_running(folder: Path) -> list[int]: + """Pids of repacks (metal_repack.py) or setup's own model preparation (repack + pack) still running for the files in + `folder`. A stale lock (its process gone) is removed. The repack and the pack move 100+ GB through the disk and one + holds up to 20 GB of memory: an engine with a 70-90 GB footprint beside them is what froze the 128 GB M5 Max twice + (watchdog panics, 2026-10-13).""" + pids = [] + if not folder.is_dir(): + return pids + for lock in [*folder.glob("*" + REPACK_LOCK), folder / PREPARE_LOCK]: + try: + pid = int(lock.read_text().strip()) + except (OSError, ValueError): + continue + if pid != os.getpid() and _pid_alive(pid): + pids.append(pid) + elif pid != os.getpid(): + try: + lock.unlink() + except OSError: + pass + return pids + + +def engine_pids(run=subprocess.run) -> list[int]: + """Pids of the running Metal engines (`strata`, the executable setup builds).""" + try: + out = run(["pgrep", "-x", "strata"], capture_output=True, text=True, timeout=10).stdout + except (OSError, subprocess.SubprocessError): + return [] + return [int(x) for x in out.split() if x.isdigit() and int(x) != os.getpid()] + + +def start(cfg_path: Path, port: int | None, open_browser: bool = True) -> int: + """serve/server.py with this config (setup.py's start without the NVIDIA / AMD card checks).""" + cfg = json.loads(cfg_path.read_text(encoding="utf-8-sig")) + missing = [p for p in [cfg["exe"], *[a for a in cfg["args"] if a.endswith(".gguf")]] if not Path(p).exists()] + if missing: + S.fail(f"{cfg_path.name} refers to missing files: {missing[0]}", "run ./setup.sh --setup to repair it") + busy = [p for g in cfg["args"] if g.endswith(".gguf") for p in repack_running(Path(g).parent)] + if busy: # never start an engine while a repack (or the pack after it) is writing: wait would hide a stuck one + S.fail(f"a model repack is still running here (pid {busy[0]}): the engine is not started beside it", + "wait for it to finish (or stop it), then start again") + cfg_path.touch() + if "--mtp" in cfg["args"][:-1]: + S.refresh_draft_vocab(Path(cfg["args"][cfg["args"].index("--mtp") + 1]), cfg.get("draft_vocab", "cjk")) + cmd = [sys.executable, str(ROOT / "serve" / "server.py"), "--engine", "strata", "--config", str(cfg_path), + "--port", str(port or cfg.get("port", 8080))] + if open_browser and cfg.get("open_browser") is not False: + cmd.append("--open") + S.say() + S.say(f" Starting {cfg.get('model_name', 'the model')} on Metal: the engine maps the model file and is ready in about " + "half a minute (longer the first time after a restart, while macOS reads the file).") + S.say(" Closing this window stops the model.") + return subprocess.call(cmd) + + +# ------------------------------------------------------------------------------------------------ main +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--family", choices=["qwen", "unsloth", "gyro"], default=None, + help="qwen = Qwen3.8-Flash-Next (GSQ-RCO), unsloth = Unsloth's UD-*_K_XL, gyro = agentionai's Gyro-S / Gyro-M " + "(follows --model)") + ap.add_argument("--model", choices=list(MAC_MODELS)) + ap.add_argument("--context", type=int) + ap.add_argument("--port", type=int) + ap.add_argument("--host") + ap.add_argument("--api-key") + ap.add_argument("--draft-vocab", choices=list(S.DRAFT_VOCABS)) + ap.add_argument("--no-browser", dest="browser", action="store_false", default=None) + ap.add_argument("--data-dir") + ap.add_argument("--move-data", action="store_true", help="with --data-dir: move the model files already here into the new folder") + ap.add_argument("--set-default-data-dir", action="store_true", help="with --data-dir: remember the new folder without moving anything") + ap.add_argument("--models-dir") + ap.add_argument("--gguf-dir") + ap.add_argument("--yes", action="store_true") + ap.add_argument("--source", choices=S.SOURCES, default=None, + help="where model files come from: auto (huggingface), huggingface, modelscope; " + "same as STRATA_SOURCE") + ap.add_argument("--setup", action="store_true") + ap.add_argument("--no-start", action="store_true") + ap.add_argument("--check", action="store_true") + ap.add_argument("--update", action="store_true") + ap.add_argument("--build", action="store_true", help="build the engine again even when it is current") + ap.add_argument("--vision", choices=["yes", "no", "gpu", "cpu"], + help="let the model read images (the image encoder, built here; yes = gpu: it reads on the Mac's GPU, cpu is ~100x " + "slower; default: asked, yes)") + a, rest = ap.parse_known_args(argv) + if rest: + S.warn("not used on a Mac: " + " ".join(rest)) + # setup.py hands a Mac over before its own parser runs, so --source has to be honoured here too: + # source_hint() recommends it on every platform, and a Mac that dropped it failed the download + # with the same unreachable host (#2363). + if a.source: + os.environ["STRATA_SOURCE"] = a.source + S.say("Strata on a Mac (Apple Silicon, Metal)") + data, elsewhere = S.data_folder_for(a) + roots = [data, *elsewhere] + models_root = Path(a.models_dir) if a.models_dir else data / "models" + + have = mac_configs() + if have and not (a.setup or a.model or a.check or a.no_start or a.update or a.build): + return start(have[0], a.port, a.browser is not False) + + # ---- 1. the Mac + S.step(1, "checking your Mac") + info = mac_info() + S.ok(f"{info['chip']}, {info['ram_gb']:.0f} GB unified memory, macOS {info['macos'] or '?'}") + problems = mac_problems(info) + for p in problems: + S.warn(p) + free = S.free_gb(data if data.exists() else ROOT) + S.ok(f"free disk: {free:.0f} GB at {data}") + if a.check: + return 1 if problems else 0 + if problems: + S.fail(problems[0]) + + # ---- 2. the choices + S.step(2, "your choices") + rec, notes = recommend_model(info["ram_gb"]) + for n in notes: + S.warn(n) + if a.model: + model = a.model + else: + names = [m for m in MAC_MODELS if not MAC_MODELS[m].get("missing")] # --model offers the others + for i, m in enumerate(names, 1): + S.say(f" {i}) {m:8} {MAC_ABOUT[m]} (maps {MAC_MODELS[m]['shard1_gb']:.0f} GB" + + ("" if MAC_MODELS[m]["tested"] else ", not measured on a Mac") + ")" + + (" (recommended)" if m in RECOMMENDED else "") + (" <- default" if m == rec else "")) + model = names[int(S.ask("Which size?", [str(i) for i in range(1, len(names) + 1)], + str(names.index(rec) + 1), a.yes)) - 1] + for n in model_note(model, info["ram_gb"]): + S.warn(n) + vision = a.vision or ("yes" if S.ask("Should the model also read images? (the image encoder, ~1 GB more)", + ["y", "n"], "y", a.yes) == "y" else "no") # on by default (strata_control/vision_reader.py) + if vision == "cpu": + from strata_control import vision_reader as VR + S.warn(VR.MAC_CPU_NOTE) + vision_gpu = vision != "cpu" + vision = "no" if vision == "no" else "yes" + ctx = a.context or recommend_context(info["ram_gb"]) + if ctx > TESTED_CONTEXT: + S.say(f" context {ctx}: measured on the Mac up to {TESTED_CONTEXT} (needles found at 8K-32K); longer is kept") + S.ok(f"{model}, context {ctx}") + family = MAC_MODELS[model].get("family", "qwen") + if a.family and a.family != family: + S.warn(f"--family {a.family}: {model} is a {family} file; using {family}") + fam = S.FAMILIES[family] + tag = fam["tag"] + model + need = (0 if a.gguf_dir else S.MODELS[model]["download_gb"]) + 6 # --gguf-dir: the files are there already + if free < need: + S.warn(f"about {need:.0f} GB are needed on the disk ({free:.0f} GB free): kept as you chose; it may run out") + + # ---- 3. Python packages (cmake and ninja come with the build in step 4, from requirements-build.txt) + S.step(3, "Python packages") + S.pip_install(S.requirement_lines() if S.REQUIREMENTS.exists() else S.PY_PACKAGES, "numpy, jinja2, ...") + + # ---- 4. the engine + S.step(4, "the Strata engine (Metal)") + llama = S.get_llama_cpp() + S.ok(f"llama.cpp {S.LLAMA_CPP_COMMIT[:7]} (ggml, gguf-py)") + if a.build: + (BUILD_DIR / "BUILD.json").unlink(missing_ok=True) + exe = build_engine(llama) + if a.update: + # update.sh's own text: "The model files are not touched". The questions are asked anyway, so when the + # answers name a model (or an image encoder) that is not here, say it: a silent exit reads like a broken + # download (an M1 Max chose IQ3_S in update.sh and got no model and no message). + files = [Path(a.gguf_dir) / n if a.gguf_dir else models_root / tag / n for n in model_shard_names(model)] + pack = S.find_in(roots, f"packs/{tag.lower()}") or data / "packs" / tag.lower() + pack_done = (pack / "native_experts.txt").exists() and (pack / "tokenizer" / "vocab.json").exists() + if not (all(f.exists() and S.done(f) for f in files) and pack_done): + S.warn(f"{model} is not installed here, and an update does not download model files: nothing was installed") + S.say(f" to download and prepare {model}: ./setup.sh --setup") + elif vision == "yes" and not (models_root / fam["mmproj"]).exists(): + S.warn("the image encoder is not downloaded here, and an update does not download model files") + S.say(" to add image support: ./setup.sh --setup") + return 0 + + # ---- 5. the model + S.step(5, f"downloading {fam['title']} {model}") + models_dir = Path(a.gguf_dir) if a.gguf_dir else models_root / tag + models_dir.mkdir(parents=True, exist_ok=True) + files = [models_dir / n for n in model_shard_names(model)] + for s in files: + if s.exists() and S.done(s): + S.ok(f"{s.name} already downloaded") + continue + S.download(fam["hf"].format(q=model) + s.name, s) + S.check_shards(files) + for s in files: # the Unsloth files: pinned sizes and SHA-256 + if shard_sha(s.name): + S.verify_sha256(s, *shard_sha(s.name)) + S.ok("model files present") + + mmproj = vexe = None + if vision == "yes": + mmproj = models_root / fam["mmproj"] + if not (mmproj.exists() and S.done(mmproj)): + S.download(fam["mmproj_hf"] + fam["mmproj"], mmproj, "vision encoder") + sha = VISION_SHA256.get(fam["mmproj"]) + if sha: + S.verify_sha256(mmproj, *sha) + vexe = build_vision(llama) + + # ---- 6. the pack and the MTP draft layer + S.step(6, "preparing the model for Strata") + pack = S.find_in(roots, f"packs/{tag.lower()}") or data / "packs" / tag.lower() + env = dict(os.environ, STRATA_GGUF_PY=str(llama / "gguf-py")) + # The repack and the pack move 100+ GB through the disk and the pack's python holds up to 20 GB: no engine may run beside + # them (an engine of 70-90 GB with a repack froze the 128 GB M5 Max twice). A lock beside the GPU's file tells the engine + # and start() to wait; a running engine stops setup here. + lock_dir = data / "metal" if len(files) > 2 else files[0].parent + native0 = data / "metal" / f"{tag.lower()}-metal.gguf" + work = (len(files) > 2 and not (native0.exists() and S.done(native0))) or \ + not ((pack / "native_experts.txt").exists() and (pack / "tokenizer" / "vocab.json").exists()) + prep_lock = lock_dir / PREPARE_LOCK + if work: + running = engine_pids() + if running: + S.fail(f"a Strata engine is running (pid {running[0]}): preparing a model beside it can freeze this Mac", + "stop the server (close its window), then run ./setup.sh --setup again") + lock_dir.mkdir(parents=True, exist_ok=True) + prep_lock.write_text(str(os.getpid())) + atexit.register(lambda: prep_lock.unlink() if prep_lock.exists() else None) + if len(files) > 2: + # many shards: the experts and dense tensors into one GGUF the GPU maps (or streams), the n-gram table read + # from its own shard (every tensor checked byte for byte: --verify) + native = data / "metal" / f"{tag.lower()}-metal.gguf" + native.parent.mkdir(parents=True, exist_ok=True) + if not (native.exists() and S.done(native)): + S.say(f" Writing the model's experts and weights into one file for the GPU (one time, " + f"{MAC_MODELS[model]['shard1_gb']:.0f} GB) ...") + S.run([sys.executable, str(ROOT / "metal" / "tools" / "metal_repack.py"), "--shard1", str(files[0]), + "--out", str(native), "--verify"]) + S.mark(native) + ple = ple_shard(files) + shards = [native, ple] + elif family == "gyro": + shards = [files[0], files[0]] # one file: the model and its n-gram table (--native, --ple-gguf) + else: + shards = files + if not (pack / "native_experts.txt").exists() or not (pack / "tokenizer" / "vocab.json").exists(): + S.run([sys.executable, str(ROOT / "tools" / "iq_pack.py"), "--gguf", str(shards[0]), "--out", str(pack), + *fam.get("pack_args", [])], env=env) # Unsloth: --compat-bf16 (the host's copies; the GPU reads the file) + S.ok(f"model prepared: {pack}") + prep_lock.unlink(missing_ok=True) + mtp = (S.find_in(roots, "mtp/mtp-q2_0.gguf") or data / "mtp" / "mtp-q2_0.gguf").parent + rt = mtp / "rt" + stream = streams(model, info["ram_gb"]) + use_mtp = stream_mtp(stream, info["ram_gb"]) + if stream and not use_mtp: + S.ok("the experts will stream from the SSD: no MTP draft layer (one token per step is faster there); " + "./setup.sh --setup on a Mac with more memory adds it") + elif stream: + S.ok(f"the experts will stream from the SSD; the MTP draft layer stays on, with --spec {STREAM_SPEC} " + "(a longer window pays when every window has a fixed cost per layer)") + if stream and info["ram_gb"] < STREAM_MTP_FROM_GB: + experts_file(shards[0]) # a small Mac: its misses are many; a big Mac's arena holds most experts + if use_mtp and not (mtp / "mtp-q2_0.gguf").exists(): + S.say(" The MTP draft layer (speculative decoding) comes from the original Qwen checkpoint: only its ~5 GB of") + S.say(" MTP tensors are downloaded.") + S.run([sys.executable, str(ROOT / "tools" / "mtp_fetch.py"), "fetch", "--out", str(mtp)], + env={**env, "STRATA_SOURCE": S.model_source()}) + S.run([sys.executable, str(ROOT / "tools" / "mtp_pack.py"), "--src", str(mtp), "--experts", "q2_0", + "--out", str(mtp / "mtp-q2_0.gguf")], env=env) + if use_mtp: + rt.mkdir(parents=True, exist_ok=True) # the Metal engine reads mtp-q2_0.gguf beside it (no rt pack) + S.refresh_draft_vocab(rt, a.draft_vocab or "cjk") + S.ok(f"MTP draft layer: {mtp / 'mtp-q2_0.gguf'}") + + # ---- 7. the config, the start script + S.step(7, "writing the start script") + port = a.port or 8080 + cfg = mac_config(exe, engine_args(pack, shards, rt, ctx, stream, info["ram_gb"]), pack, f"{fam['name']}-{model.lower()}", tag, port) + cfg["env"] = {"STRATA_METAL_STREAM": "1" if stream else "0"} # the engine follows setup's measured choice + if vision == "yes": + cfg["args"] += ["--vision"] + cfg["vision"] = {"exe": str(vexe), "mmproj": str(mmproj), "model": str(files[0]), "gpu": vision_gpu, + "max_tokens": 1024} + if a.host: + cfg["host"] = a.host + if a.api_key: + cfg["api_key"] = a.api_key + if a.draft_vocab: + cfg["draft_vocab"] = a.draft_vocab + if a.browser is not None: + cfg["open_browser"] = a.browser + cfg_path = ROOT / f"strata-{tag.lower()}.json" + S.write_config(cfg_path, cfg) + script = S.write_run_script(tag, cfg_path, port, cfg.get("open_browser") is not False) + S.ok(f"start script: {script.name}") + S.say() + S.say("All set.") + S.say(f" API (OpenAI): http://127.0.0.1:{port}/v1") + S.say(f" API (Anthropic): http://127.0.0.1:{port}/v1/messages") + S.say(f" Next time: just run ./setup.sh (or {script.name}) - it starts right away") + if a.no_start: + return 0 + return start(cfg_path, port, a.browser is not False) + + +if __name__ == "__main__": + try: + sys.exit(main()) + except KeyboardInterrupt: + S.say("\nstopped.") + sys.exit(1) diff --git a/edge0-strata/metal/src/engine/arena_budget.h b/edge0-strata/metal/src/engine/arena_budget.h new file mode 100644 index 00000000..4cc024b7 --- /dev/null +++ b/edge0-strata/metal/src/engine/arena_budget.h @@ -0,0 +1,177 @@ +// metal/src/engine/arena_budget.h - how big the SSD-streaming expert arena may be (stream.h), and how much the engine may +// hold resident in all, as plain arithmetic so it can be tested without a Mac (metal/tests/arena_budget_test.cpp). +// +// The arena is anonymous memory the GPU needs resident while a command buffer runs. It must leave room for macOS and the +// desktop, and for the other big memory the engine itself uses while it runs: +// keep macOS and the desktop: max(7 GB, 15% of the RAM) - STRATA_METAL_KEEP_FREE_GB sets it. The M5 Max +// (128 GB) froze with an 89 GB arena and 7 GB kept beside a 50 GB n-gram table; with 15% it keeps 20.6 GB. +// The 32 GB M4 (IQ3_XXS after a 6K prompt, 5 interleaved runs, memory_pressure every 3 s, vm_stat +// swapouts): 7 GB kept + 2 GB reserve 5.93 tok/s (arena ~8.5 GB, least free 50%), 7 + 0.5 6.59 (~10 GB, +// 49%), 5 + 0.5 7.47 (~12 GB, 43%); no swap-out in any run. But a 34-minute server soak with 5 GB kept +// swapped out 93 MB while its 11.7 GB arena filled (none after): 7 GB again. +// ple the n-gram table: mapped whole in RAM (--ple-io ram) it is all of it; read row by row (the default, +// F_NOCACHE: no file cache) only its bounded row cache, 1M rows of at most 192 bytes: 0.5 GB reserved. +// GPU the Metal working set less what is already allocated, less 1 GB. +// footprint the engine's WHOLE anonymous footprint - arena + every Metal buffer + what the process holds beside them +// (host memory) - stays under footprint_cap() (below): on a big Mac 60% of the RAM. The 67% rule it replaces +// counted only the Metal buffers (cx.allocated(), 7.5 GB) while the process held 11-15 GB, and 67% was too +// much anyway (see kLeaveShare). +// STRATA_METAL_ARENA_GB sets the arena outright (the engine says when that goes over the cap; it is the user's choice). +// The arena is placed just before it is filled, and the memory macOS can give is asked again then: when it has dropped +// (another program), the arena shrinks (arena_shrink). +// "Available" (arena_avail) counts free and purgeable pages and the file cache macOS drops without writing anything: +// inactive pages, but no more of them than there are file-backed pages. Inactive anonymous memory (other programs') +// is left out: macOS frees it only by compressing or swapping it (the soak's 93 MB swapped out while an arena planned +// with every inactive page filled). Free + purgeable alone is too little: at the engine's start the model file's own +// pages fill the cache (6.4 GB free + purgeable on the 32 GB M4 against 15 GB a minute later). The arena is filled in +// steps (STRATA_METAL_ARENA_STEP_GB, default 1.5): before each step the free memory is asked again, and the fill stops - +// the rest of the arena's slots unused - when a step would take macOS below what it keeps (arena_grow). +#pragma once + +#include +#include + +namespace strata::metal_engine { + +struct ArenaMem { + double phys = 0; // bytes of RAM + double working_set = 0; // the Metal device's recommended working set + double allocated = 0; // already allocated by the engine (Metal buffers) + double available = 0; // what macOS can give now without paging anything out (arena_avail) + double ple_mapped = 0; // the n-gram table's bytes when it is mapped whole in RAM, else 0 + double proc_footprint = 0; // what the process holds right now, whatever it is (task_info phys_footprint); 0: unknown + double filecache_min = 0; // sysctl vm.vm_page_filecache_min in bytes: the file cache macOS keeps before it compresses (0: unknown) + double keep_gb = -1; // STRATA_METAL_KEEP_FREE_GB (< 0: unset) + double arena_gb = -1; // STRATA_METAL_ARENA_GB (< 0: unset) + double leave_share = -1; // STRATA_METAL_LEAVE_SHARE: the share of the RAM a big Mac leaves to macOS (< 0: kLeaveShare) +}; + +constexpr double kGB = 1e9; + +// Why a cap at all. macOS keeps a floor of file cache (vm.vm_page_filecache_min: 37.2 GB, 27% of the 128 GB M5 Max) and, +// when anonymous memory wants more than the rest, compresses the anonymous pages (then swaps) instead of dropping file +// cache below it. The arena is anonymous and the GPU wires it only while a command buffer runs: while it is being filled, +// or idle, it is ordinary pageable memory. Measured on the M5 Max (UD-Q6_K_XL streaming, memwatch every second): +// the old build, a 100 GB arena (what it planned after a fresh boot): 19 GB of the arena compressed in 7 s, 29,000 pages +// swapped out, pressure "warn" - the signature of the two panics (compressor at 3.9M pages, free memory 13 MB, the +// engine 95.3 and 89.9 GiB resident, watchdogd starved for 92 s); +// the old build, auto arena 69 GB, a 40 GB file write and a 16 GB program beside it: compressor 21k -> 4.07M pages in +// 10 s, swapouts +29,000, pressure warn (stopped by the watchdog); +// footprint cap at 70% of the RAM (wired 94.5 GB), 40 GB file write: compressor 1.8M pages, swapouts +12,000, stopped; +// 35% left (89 GB): clear with the write alone (compressor 0.28M pages, 316 pages swapped), not with more beside it; +// 40% left (cap 82.5 GB, footprint 77 GB): 40 GB write + 16 GB program, swapouts 0, compressor 0.3M pages, decode and +// prompt speed equal to the 80 GB arena's (51 tok/s story, 1,020 tok/s for 15K). +// And the arena is wired (mlock, stream.cpp) once it is filled: unwired it is pageable whenever the GPU is idle (the GPU's own wiring +// ends seconds after its last command buffer; a Metal residency set did not hold it either), and a 20 GB file written beside the +// idle engine took tens of GB of the arena into the compressor and froze the Mac (a third watchdog panic, in a control run of this +// work: wired 2.4 GB, file cache 50 GB, compressor 1.9M pages, 45M compressions). Wired, the same write (full speed, 20 GB in 2.5 s) +// raised the compressor by 0.7M pages (the ~11 GB of the engine that is not arena), swapped out 408 pages and changed nothing else. +// A pack that stays RESIDENT is different: its bytes are clean file pages (macOS drops them, never compresses them) and they +// count as file cache, not against the floor. UD-Q4_K_XL resident on the M5 Max (wired 91.7 GB, 67%), 40 GB write + 16 GB +// program beside it: swapouts 0, compressor 5,500 pages, 63 tok/s. So two caps: +// footprint_cap anonymous memory (the streaming arena, Metal buffers, host memory): kLeaveShare = 40% of the RAM left +// on a big Mac (at least the file-cache floor + 12 GB) +// resident_cap a resident pack's mapped bytes: kResidentLeaveShare = 30% left (the old rule, measured as above) +constexpr double kLeaveShare = 0.40; +constexpr double kResidentLeaveShare = 0.30; +constexpr double kBigMac = 64.0 * kGB; +constexpr double kFileCacheHeadroom = 12.0 * kGB; // beside the file-cache floor: macOS itself (~6 GB) and one more program + +inline double arena_keep(const ArenaMem& m) { + return m.keep_gb >= 0 ? m.keep_gb * kGB : std::max(7.0 * kGB, 0.15 * m.phys); +} + +/// The most the engine may hold as anonymous memory, all counted (arena + Metal buffers + host memory). +inline double footprint_cap(const ArenaMem& m) { + const double share = m.leave_share >= 0 ? m.leave_share : kLeaveShare; + double leave = arena_keep(m); + if (m.phys >= kBigMac) leave = std::max({leave, share * m.phys, m.filecache_min > 0 ? m.filecache_min + kFileCacheHeadroom : 0.0}); + return std::min(m.working_set, m.phys - leave); +} + +/// The most a resident pack (its mapped bytes + the engine's buffers) may hold. +inline double resident_cap(const ArenaMem& m) { + double leave = arena_keep(m); + if (m.phys >= kBigMac) leave = std::max(leave, kResidentLeaveShare * m.phys); + return std::min(m.working_set, m.phys - leave); +} + +/// What the engine holds before the arena: the Metal buffers (alloc + wrap) or, when the process holds more, all of it +/// (host memory, mapped tables). +inline double footprint_base(const ArenaMem& m) { return std::max(m.allocated, m.proc_footprint); } + +/// Bytes macOS can give without compressing or swapping: free + purgeable + the clean file cache (inactive and +/// speculative pages, at most the file-backed ones). +inline double arena_avail(double free_b, double purgeable_b, double inactive_b, double speculative_b, double file_b) { + return free_b + purgeable_b + std::min(inactive_b + speculative_b, file_b); +} + +inline double arena_ple_reserve(const ArenaMem& m) { return m.ple_mapped > 0 ? m.ple_mapped : 0.5 * kGB; } + +/// The arena's bytes (at least 1 GB: a smaller arena still runs, at a low hit rate). +inline double arena_budget(const ArenaMem& m) { + if (m.arena_gb >= 0) return m.arena_gb * kGB; + const double gpu = m.working_set - footprint_base(m) - 1.0 * kGB; + const double ram = m.available - arena_keep(m) - arena_ple_reserve(m); + const double foot = footprint_cap(m) - footprint_base(m) - arena_ple_reserve(m); // the whole footprint under the cap + return std::max(1.0 * kGB, std::min(std::min(gpu, ram), foot)); +} + +/// The plan as numbers, for the log line and the tests: what the engine will hold with `arena` bytes of arena. +struct FootprintPlan { + double arena = 0, base = 0, ple = 0, total = 0, cap = 0; + bool over() const { return total > cap + 1e6; } +}; +inline FootprintPlan footprint_plan(const ArenaMem& m, double arena) { + FootprintPlan p; + p.arena = arena; + p.base = footprint_base(m); + p.ple = arena_ple_reserve(m); + p.total = p.arena + p.base + p.ple; + p.cap = footprint_cap(m); + return p; +} + +/// A pack that stays resident (no SSD streaming) holds its mapped bytes plus `overhead` (the prompt buffers, KV, the draft +/// layer, ...): does that fit under the resident cap? (stream.cpp: ExpertStream::wanted) +inline bool resident_fits(double phys, double working_set, double model_bytes, double overhead) { + ArenaMem m; + m.phys = phys; + m.working_set = working_set; + return model_bytes + overhead <= resident_cap(m); +} + +/// Just before the arena is filled: the bytes it may still have when macOS can now give `available_now` (it shrinks, +/// never grows; the explicit STRATA_METAL_ARENA_GB is kept as it is). +inline double arena_shrink(const ArenaMem& m, double planned, double available_now) { + if (m.arena_gb >= 0) return planned; + ArenaMem now = m; + now.available = available_now; + return std::min(planned, arena_budget(now)); +} + +/// The stepped fill: may the next `step` bytes of the arena be filled while macOS has `available_now` free + purgeable? +/// (the explicit STRATA_METAL_ARENA_GB always fills) +inline bool arena_grow(const ArenaMem& m, double available_now, double step) { + if (m.arena_gb >= 0) return true; + return available_now - step >= arena_keep(m) + arena_ple_reserve(m); +} + +/// The fill as a plan, for the test: the pairs' bytes in rank order, steps of `step` bytes, `avail(filled)` the free +/// memory once `filled` bytes are in. Returns how many pairs are filled. +template +inline size_t arena_fill_plan(const ArenaMem& m, const double* cost, size_t n, double step, Avail avail) { + size_t i = 0; + double filled = 0; + while (i < n) { + double s = 0; + size_t j = i; + while (j < n && (s < step || j == i)) s += cost[j++]; + if (!arena_grow(m, avail(filled), s)) break; + filled += s; + i = j; + } + return i; +} + +} // namespace strata::metal_engine diff --git a/edge0-strata/metal/src/engine/ds4_ctx.cpp b/edge0-strata/metal/src/engine/ds4_ctx.cpp new file mode 100644 index 00000000..5d7984d4 --- /dev/null +++ b/edge0-strata/metal/src/engine/ds4_ctx.cpp @@ -0,0 +1,80 @@ +// metal/src/engine/ds4_ctx.cpp - the one-Context allocator the dsv41 engine can opt into. +// +// Why: a Metal kernel must be handed the gpuAddress of a buffer that the SAME mtl::Context encodes the dispatch for (its +// residency set). strata::gpu, which metal_compat's cudaMalloc maps to, is a CUDA-shaped facade with its own device state, so a +// dsv41 launch would read a foreign residency set. This file gives metal_compat an alternative: allocate inside one Context and +// remember gpuAddress -> Buffer, so wrap()/Arg::B can adopt the memory without a copy. +// +// Off by default: the engine's default path stays exactly as it was. Set STRATA_DS4_METAL_CTX=1 to route allocations here - +// which becomes testable the moment a pack exists (the engine allocates only after it has read ds4.txt). +#include "mtl_rt.h" + +#include +#include +#include +#include + +namespace strata::dsv41::metal_ctx { + +namespace { +struct Entry { + mtl::Buffer buf; + size_t bytes = 0; +}; + +std::mutex g_mu; +std::unordered_map* g_map = nullptr; // gpuAddress -> buffer +mtl::Context* g_cx = nullptr; + +bool enabled() { + static const bool on = std::getenv("STRATA_DS4_METAL_CTX") != nullptr; + return on; +} + +mtl::Context& cx() { + if (!g_cx) g_cx = new mtl::Context(); + return *g_cx; +} +} // namespace + +bool active() { return enabled(); } + +void* alloc(size_t bytes) { + if (!enabled()) return nullptr; + const size_t page = 16384; + const size_t rounded = (bytes + page - 1) / page * page; + mtl::Buffer b = cx().alloc(rounded); + if (!b.handle || !b.ptr) { + std::fprintf(stderr, "dsv41 metal: Context::alloc(%zu) failed\n", bytes); + return nullptr; + } + std::lock_guard lk(g_mu); + if (!g_map) g_map = new std::unordered_map(); + (*g_map)[b.ptr] = Entry{b, rounded}; + return b.ptr; // the CPU address doubles as the handle's key; the device pointer is b.handle's gpuAddress +} + +void* host_view_of(void* p) { + std::lock_guard lk(g_mu); + if (!g_map) return nullptr; + auto it = g_map->find(p); + return it == g_map->end() ? nullptr : it->second.buf.ptr; +} + +const mtl::Buffer* buffer_of(void* p) { + std::lock_guard lk(g_mu); + if (!g_map) return nullptr; + auto it = g_map->find(p); + return it == g_map->end() ? nullptr : &it->second.buf; +} + +void release(void* p) { + std::lock_guard lk(g_mu); + if (!g_map) return; + auto it = g_map->find(p); + if (it == g_map->end()) return; + cx().free(it->second.buf); + g_map->erase(it); +} + +} // namespace strata::dsv41::metal_ctx diff --git a/edge0-strata/metal/src/engine/ds4_finish.cpp b/edge0-strata/metal/src/engine/ds4_finish.cpp new file mode 100644 index 00000000..d755e911 --- /dev/null +++ b/edge0-strata/metal/src/engine/ds4_finish.cpp @@ -0,0 +1,2754 @@ +// metal/src/engine/ds4_finish.cpp - the host side of ds4_finish (ds4_kernels.cu:1894, finish_kernel). +// +// The kernel takes the engine's own device pointer for x as a bound buffer plus an offset, and the Ds4Finish struct as a constant +// buffer whose pointer fields are the engine's gpuAddresses (Metal lets a kernel dereference them - ss the repo's note that +// "kernels may take raw gpuAddress pointers"). Strides and counts travel as plain values. +// +// The mrope (rd > 0) and rounding (quant != 0) branches are not written yet, and neither is the KV pack (pdst): this function +// returns false for those, so the engine reports the launch line instead of a step being skipped. +#include "mtl_rt.h" + +#include +#include + +#include +#include +#include +#include +#include +#include + +// called by the gpu layer's synchronisation: everything batched must reach the GPU before a host read +extern "C" void ds4_metal_flush(); +extern "C" size_t ds4_gpu_alloc_remaining(const void* cpu); +extern "C" void* ds4_gpu_buffer_for_cpu(const void* cpu, size_t bytes, size_t* off); +extern "C" void ds4_gpu_sync_all_streams(); +extern "C" unsigned long long ds4_dispatch_seq(); +static unsigned long long g_q8_deduped = 0; + +// The bytes a kernel's wrap() binds from the page-aligned `base`, `need` of them used: whole pages, but never past the end of +// the allocation that holds them. A range that leaves its allocation cannot be bound as the runtime's own buffer and gets a fresh +// no-copy buffer instead: a second GPU mapping of the same pages, through which a kernel read the previous values of what an +// earlier kernel had written through the first (ext_to_ring read the window rows of the layer before; a prompt of a few tokens +// answered differently on every run, with the same result only when the GPU was synchronised after the kernel). +static size_t wrap_len(size_t base, size_t need) { + const size_t pg = 16384, up = (need + pg - 1) / pg * pg; + const size_t rem = ds4_gpu_alloc_remaining((const void*) base); + return rem > 0 && rem < up ? rem : up; // (sparse_attn asks for one row more than the window buffer holds, which it never reads) +} + +#define DSV41_PROBE(fmt, ...) do { static bool once_ = false; if (!once_) { once_ = true; std::fprintf(stderr, "[probe] " fmt, __VA_ARGS__); } } while (0) + +namespace strata::dsv41::metal_host { + +// mirrors metal/src/engine/msl/ds4_init.metal's Ds4Finish (Metal has no 'long long': 64-bit is 'long') +struct GpuFinish { + uint64_t w; // device float* (the engine's gpuAddress, 0 = none) + float eps; + int rd, pos, inverse, vpp, pos_step; + uint64_t freqs; // device float* + int quant; + uint64_t dst; // device ushort* + int64_t dst_stride; + uint64_t pdst; // device uchar* + int64_t pdst_stride; +}; + +namespace { +mtl::Context* g_cx = nullptr; +int g_lib = -1; +int g_lib_iq = -1; + +bool ensure() { + if (std::getenv("DS4_DBG_MEM")) { // does a write through host_view read back through cudaMemcpy? + void* p = nullptr; + const size_t tsz = std::getenv("DS4_MEM_GB") ? 64u << 20 : 4096u; + if (cudaMalloc(&p, tsz) == cudaSuccess) { + void* v = strata::gpu::host_view(p); + if (v) { + ((float*) v)[0] = 42.0f; + float out = 0.0f; + cudaMemcpy(&out, p, 4, cudaMemcpyDeviceToHost); + std::fprintf(stderr, "DBG mem: gpu=%p view=%p wrote 42 through the view, cudaMemcpy gives %g\n", p, v, out); + } else { + std::fprintf(stderr, "DBG mem: gpu=%p has no host view\n", p); + } + cudaFree(p); + } + } + if (g_cx) return g_lib >= 0; + g_cx = new mtl::Context(); + g_cx->start_keepalive(); + g_lib = g_cx->compile({"../src/engine/msl/ds4_init.metal"}, "", false); + if (g_lib < 0) std::fprintf(stderr, "dsv41 metal: compiling ds4_init.metal failed\n"); + g_lib_iq = g_cx->compile({"../src/engine/msl/ds4_iq.metal"}, "", false); + if (g_lib_iq < 0) std::fprintf(stderr, "dsv41 metal: compiling ds4_iq.metal failed\n"); + if (g_lib >= 0 && std::getenv("DS4_DBG_BIND")) { // write a constant by each route, copy it back + void* p = nullptr; + const size_t tsz = std::getenv("DS4_MEM_GB") ? 64u << 20 : 4096u; + if (cudaMalloc(&p, tsz) == cudaSuccess) { + void* v = strata::gpu::host_view(p); + float got[4] = {0, 0, 0, 0}; + { + mtl::Buffer b = g_cx->wrap((void*) (((uintptr_t) v) & ~(uintptr_t) 16383), 16384); + mtl::Pipeline pl = g_cx->pipeline(g_lib, "ds4_test_write"); + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(b, ((uintptr_t) v) & 16383), mtl::Arg::V((unsigned) 4)}; + d.grid[0] = 1; d.grid[1] = d.grid[2] = 1; d.tg[0] = 32; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + g_cx->flush_runs(); + cudaMemcpy(got, p, sizeof(got), cudaMemcpyDeviceToHost); + std::fprintf(stderr, "DBG bind A (wrap host_view + Arg::B): %g %g %g %g\n", got[0], got[1], got[2], got[3]); + } + { + uint64_t pp = (uint64_t) (uintptr_t) p; + mtl::Pipeline pl = g_cx->pipeline(g_lib, "ds4_test_write_v"); + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::V(pp), mtl::Arg::V((unsigned) 4)}; + d.grid[0] = 1; d.grid[1] = d.grid[2] = 1; d.tg[0] = 32; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + g_cx->flush_runs(); + cudaMemcpy(got, p, sizeof(got), cudaMemcpyDeviceToHost); + std::fprintf(stderr, "DBG bind B (device pointer by value): %g %g %g %g\n", got[0], got[1], got[2], got[3]); + } + cudaFree(p); + } + } + return g_lib >= 0; +} +} // namespace + +// the engine's Finish, as ds4_kernels.cuh declares it +struct FinishIn { + const float* w; + float eps; + int rd, pos, inverse, vpp, pos_step; + const float* freqs; + int quant; + const uint16_t* dst; + long long dst_stride; + uint8_t* pdst; + long long pdst_stride; +}; + +bool finish(float* x, int n_vec, int len, long long stride, const FinishIn& f) { + if (n_vec <= 0 || len <= 0) return true; + // finish_kernel (the KV pack: the row quantised to codes + group scales) or ; the rounding + // branch (quant without pdst) is the one still unwritten + const bool pack = f.pdst != nullptr; + if (!ensure()) return false; + void* hx = strata::gpu::host_view(x); + if (!hx) { std::fprintf(stderr, "dsv41 metal: finish has no host view\n"); return false; } + const size_t page = 16384; + const size_t bytes = (size_t) (n_vec - 1) * (size_t) stride * sizeof(float) + (size_t) len * sizeof(float); + const size_t base = ((uintptr_t) hx) & ~(uintptr_t) (page - 1); + const size_t off = ((uintptr_t) hx) - base; + mtl::Buffer bx = g_cx->wrap((void*) base, wrap_len(base, bytes + off)); + if (!bx.handle) { std::fprintf(stderr, "dsv41 metal: finish wrap failed\n"); return false; } + + GpuFinish gf{}; + gf.w = (uint64_t) (uintptr_t) f.w; + gf.eps = f.eps; + gf.rd = f.rd; gf.pos = f.pos; gf.inverse = f.inverse; gf.vpp = f.vpp; gf.pos_step = f.pos_step; + gf.freqs = (uint64_t) (uintptr_t) f.freqs; + gf.quant = f.quant; + gf.dst = (uint64_t) (uintptr_t) f.dst; + gf.dst_stride = f.dst_stride; + gf.pdst = (uint64_t) (uintptr_t) f.pdst; + gf.pdst_stride = f.pdst_stride; + + mtl::Pipeline pl; + try { + pl = g_cx->pipeline(g_lib, pack ? "ds4_finish_pack" : "ds4_finish"); + } catch (...) { + return false; + } + // one dispatch, one threadgroup a vector (the kernel's blockIdx.x is threadgroup_position_in_grid) + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bx, off), mtl::Arg::V((long) stride), mtl::Arg::V((unsigned) len), mtl::Arg::V(gf)}; + d.grid[0] = (uint32_t) n_vec; + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 256; + d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + // The fp8-rounding finishes (the KV rows, the attention output) are read stale by the next dispatch when it shares their + // command buffer: end the buffer after them (committed, not waited for: the queue orders the buffers). + // (it once ended the command buffer after every fp8-rounding finish, 2 a layer, against a stale read; each extra buffer costs the GPU ~0.3 ms idle, a third of a decode token. + // The decode-path score is bit-identical without it, the dispatch barrier after every dispatch orders the readers. DS4_FINISH_SPLIT=1 brings it back.) + static const bool split_after = std::getenv("DS4_FINISH_SPLIT") != nullptr; + if (f.quant && split_after) g_cx->split_runs(); + if (const char* fm2 = std::getenv("DS4_FLUSH_FINISH_CTX")) { if (f.quant) { if (std::atoi(fm2) == 2) g_cx->split_runs(); else g_cx->flush_runs(); } } + if (const char* fm = std::getenv("DS4_FLUSH_FINISH")) { // debugging: which finish calls need a drain after them + const int m = std::atoi(fm); + if ((m == 1 && f.dst) || (m == 2 && f.rd > 0) || (m == 3 && f.w) || (m == 4 && f.quant) || (m == 5 && !f.dst) || + (m == 6 && f.dst == nullptr && f.rd == 0) || (m == 7 && f.rd > 0 && f.inverse) || (m == 8 && f.rd > 0 && !f.inverse && !f.dst) || (m == 9 && f.rd > 0 && !f.dst) || (m == 10 && (f.dst || f.inverse)) || (m == 11 && f.rd > 0 && !f.inverse)) + ds4_metal_flush(); + } + return true; +} + + +// ds4_kernels.cu:1677 - hc_post_inplace: h[t][k][e] = post[t][k]*y[t][e] + sum_j comb[t][j][k]*h_old[t][j][e] +bool hc_post(const float* y, float* h, const float* post, const float* comb, int dim, int hc, int T) { + if (!ensure()) return false; + void* hy = strata::gpu::host_view((void*) y); + void* hh = strata::gpu::host_view(h); + if (!hy || !hh) { std::fprintf(stderr, "dsv41 metal: hc_post has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer by = wrap(hy, (size_t) T * dim * 4 ); + mtl::Buffer bh = wrap(hh, (size_t) T * hc * dim * 4 ); + if (!by.handle || !bh.handle) return false; + // post is [T][8], comb is [T][64]: device operands like the others (reached through their host views) + void* hpost = strata::gpu::host_view((void*) post); + void* hcomb = strata::gpu::host_view((void*) comb); + if (!hpost || !hcomb) { std::fprintf(stderr, "dsv41 metal: hc_post has no host view for post / comb\n"); return false; } + mtl::Buffer bp = wrap(hpost, (size_t) T * 8 * 4); + mtl::Buffer bc = wrap(hcomb, (size_t) T * 64 * 4); + if (!bp.handle || !bc.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_hc_post"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(by, off(hy)), mtl::Arg::B(bh, off(hh)), mtl::Arg::B(bp, off(hpost)), mtl::Arg::B(bc, off(hcomb)), + mtl::Arg::V((unsigned) dim), mtl::Arg::V((unsigned) hc)}; + d.grid[0] = (uint32_t) ((dim + 255) / 256); + d.grid[1] = (uint32_t) T; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + const bool dbg = std::getenv("DS4_DBG_SG") != nullptr; + if (dbg) { + g_cx->flush_runs(); + const float* fy = (const float*) hy; const float* fh = (const float*) hh; const float* fp = (const float*) hpost; const float* fc = (const float*) hcomb; + std::fprintf(stderr, "hc_post in: y %g %g h %g %g | %g %g post %g %g %g %g comb %g %g %g %g\n", fy[0], fy[1], fh[0], fh[1], fh[(size_t) dim], fh[(size_t) dim + 1], fp[0], fp[1], fp[2], fp[3], fc[0], fc[1], fc[2], fc[3]); + } + g_cx->run({d}); + if (dbg) { + g_cx->flush_runs(); + const float* fh = (const float*) hh; + std::fprintf(stderr, "hc_post out: h %g %g | %g %g\n", fh[0], fh[1], fh[(size_t) dim], fh[(size_t) dim + 1]); + } + g_cx->free(bp); + g_cx->free(bc); + return true; +} + + +// ds4_kernels.cu:1692 - hc_mix_scale: mixes[t][i] *= rsqrt(sum(h[t])^2 / hcdim + eps) +bool hc_mix_scale(const float* h, int hcdim, int n_mix, int T, float eps, float* mixes) { + if (!ensure()) return false; + void* hh = strata::gpu::host_view((void*) h); + void* hm = strata::gpu::host_view((void*) mixes); + if (!hh || !hm) { std::fprintf(stderr, "dsv41 metal: hc_mix_scale has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bh = wrap(hh, (size_t) T * hcdim * 4 ); + mtl::Buffer bm = wrap(hm, (size_t) T * 32 * 4 ); + if (!bh.handle || !bm.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_hc_mix_scale"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bh, off(hh)), mtl::Arg::B(bm, off(hm)), mtl::Arg::V((unsigned) hcdim), + mtl::Arg::V((unsigned) n_mix), mtl::Arg::V(eps)}; + d.grid[0] = (uint32_t) T; + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + if (std::getenv("DS4_DBG_SG")) { + g_cx->flush_runs(); + const float* fc = (const float*) hm; + std::fprintf(stderr, "hc_mix_scale T %d hcdim %d n_mix %d eps %g -> %g %g %g\n", T, hcdim, n_mix, eps, fc[0], fc[1], fc[2]); + } + return true; +} + + +// ds4_kernels.cu:1901 - hc_mix: mixes[t][m] = (h[t] . fn[m]) * rsqrt(mean(h[t]^2) + eps) +bool hc_mix(const float* h, const float* fn, int hcdim, int n_mix, int T, float eps, float* mixes) { + if (!ensure()) return false; + void* hh = strata::gpu::host_view((void*) h); + void* hf = strata::gpu::host_view((void*) fn); + void* hm = strata::gpu::host_view((void*) mixes); + if (!hh || !hf || !hm) { std::fprintf(stderr, "dsv41 metal: hc_mix has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bh = wrap(hh, (size_t) T * hcdim * 4 ); + mtl::Buffer bf = wrap(hf, (size_t) n_mix * hcdim * 4 ); + mtl::Buffer bm = wrap(hm, (size_t) T * 32 * 4 ); + if (!bh.handle || !bf.handle || !bm.handle) return false; + mtl::Pipeline pl; + const bool all = T >= 16 && n_mix <= 32 && !std::getenv("DS4_NO_HC_ALL"); // a prompt chunk: a threadgroup a token, 4 mixes a pass + try { pl = g_cx->pipeline(g_lib, all ? "ds4_hc_mix_all" : "ds4_hc_mix"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bh, off(hh)), mtl::Arg::B(bf, off(hf)), mtl::Arg::B(bm, off(hm)), + mtl::Arg::V((unsigned) hcdim), mtl::Arg::V(eps)}; + if (all) d.args.push_back(mtl::Arg::V((unsigned) n_mix)); + d.grid[0] = all ? (uint32_t) T : (uint32_t) n_mix; + d.grid[1] = all ? 1u : (uint32_t) T; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + + +// ds4_kernels.cu:1908 - hc_split: the sinkhorn normalisation (scalar path; see the MSL comment) +bool hc_split(const float* mixes, const float* scale, const float* base, int hc, int iters, float hc_eps, int T, + float* pre, float* post, float* comb) { + if (!ensure()) return false; + void* hm = strata::gpu::host_view((void*) mixes); + void* hp = strata::gpu::host_view((void*) pre); + void* ho = strata::gpu::host_view((void*) post); + void* hc_ = strata::gpu::host_view((void*) comb); + if (!hm || !hp || !ho || !hc_) { std::fprintf(stderr, "dsv41 metal: hc_split has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bm = wrap(hm, (size_t) T * 32 * 4 ); + mtl::Buffer bp = wrap(hp, (size_t) T * 8 * 4 ); + mtl::Buffer bo = wrap(ho, (size_t) T * 8 * 4 ); + mtl::Buffer bc = wrap(hc_, (size_t) T * 64 * 4 ); + if (!bm.handle || !bp.handle || !bo.handle || !bc.handle) return false; + // scale and base come from the engine like every other operand: reach them through host_view, not as raw host pointers + void* hs = strata::gpu::host_view((void*) scale); + void* hb = strata::gpu::host_view((void*) base); + if (!hs || !hb) { std::fprintf(stderr, "dsv41 metal: hc_split scale/base have no host view\n"); return false; } + const size_t nbase = (size_t) (2 * hc + hc * hc); + mtl::Buffer bs = wrap(hs, 3 * sizeof(float)); + mtl::Buffer bb = wrap(hb, nbase * sizeof(float)); + if (!bs.handle || !bb.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, hc == 4 && !std::getenv("DS4_NO_HCS4") ? "ds4_hc_split4" : "ds4_hc_split"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bm, off(hm)), mtl::Arg::B(bs, off(hs)), mtl::Arg::B(bb, off(hb)), mtl::Arg::B(bp, off(hp)), + mtl::Arg::B(bo, off(ho)), mtl::Arg::B(bc, off(hc_)), mtl::Arg::V((unsigned) hc), + mtl::Arg::V((unsigned) iters), mtl::Arg::V((float) hc_eps)}; + d.grid[0] = (uint32_t) T; + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 32; d.tg[1] = d.tg[2] = 1; + if (std::getenv("DS4_DBG_SG")) { + g_cx->flush_runs(); + const float* fm = (const float*) hm; const float* fs = (const float*) hs; const float* fb = (const float*) hb; + std::fprintf(stderr, "hc_split in: mix %g %g %g scale %g %g %g base %g %g\n", fm[0], fm[1], fm[2], fs[0], fs[1], fs[2], fb[0], fb[1]); + } + g_cx->run({d}); + if (std::getenv("DS4_DBG_SG")) { + g_cx->flush_runs(); + const float* fp = (const float*) hp; + std::fprintf(stderr, "hc_split out: pre %g %g %g %g\n", fp[0], fp[1], fp[2], fp[3]); + } + return true; +} + + +// the hyper-connection mix and split in one dispatch (ds4_hc_mix_split): the engine used to issue hc_mix and hc_split +// back to back for the same token, and the kernel does both in one threadgroup, bit for bit as before. +bool hc_mix_split(const float* h, const float* fn, int hcdim, int n_mix, int T, float eps, float* mixes, + const float* scale, const float* base, int hc, int iters, float hc_eps, float* pre, float* post, + float* comb) { + if (!ensure()) return false; + void* hh = strata::gpu::host_view((void*) h); + void* hf = strata::gpu::host_view((void*) fn); + void* hm = strata::gpu::host_view((void*) mixes); + void* hs = strata::gpu::host_view((void*) scale); + void* hb = strata::gpu::host_view((void*) base); + void* hp = strata::gpu::host_view((void*) pre); + void* ho = strata::gpu::host_view((void*) post); + void* hc_ = strata::gpu::host_view((void*) comb); + if (!hh || !hf || !hm || !hs || !hb || !hp || !ho || !hc_) { + std::fprintf(stderr, "dsv41 metal: hc_mix_split has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base_ = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base_, wrap_len(base_, end - base_)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(size_t) (page - 1))); }; + mtl::Buffer bh = wrap(hh, (size_t) T * hcdim * 4); + mtl::Buffer bf = wrap(hf, (size_t) n_mix * hcdim * 4); + mtl::Buffer bm = wrap(hm, (size_t) T * 32 * 4); + mtl::Buffer bs = wrap(hs, 3 * sizeof(float)); + mtl::Buffer bb = wrap(hb, (size_t) (2 * hc + hc * hc) * sizeof(float)); + mtl::Buffer bp = wrap(hp, (size_t) T * 8 * 4); + mtl::Buffer bo = wrap(ho, (size_t) T * 8 * 4); + mtl::Buffer bc = wrap(hc_, (size_t) T * 64 * 4); + if (!bh.handle || !bf.handle || !bm.handle || !bs.handle || !bb.handle || !bp.handle || !bo.handle || !bc.handle) + return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_hc_mix_split"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bh, off(hh)), mtl::Arg::B(bf, off(hf)), mtl::Arg::B(bm, off(hm)), + mtl::Arg::V((unsigned) hcdim), mtl::Arg::V((float) eps), mtl::Arg::V((unsigned) n_mix), + mtl::Arg::B(bs, off(hs)), mtl::Arg::B(bb, off(hb)), mtl::Arg::B(bp, off(hp)), + mtl::Arg::B(bo, off(ho)), mtl::Arg::B(bc, off(hc_)), mtl::Arg::V((unsigned) hc), + mtl::Arg::V((unsigned) iters), mtl::Arg::V((float) hc_eps)}; + d.grid[0] = (uint32_t) T; + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + + +// ds4_kernels.cu:1938 - moe_sort: count each virtual slot's picks, prefix-sum the counts, fill in token order. +// Three dispatches on one command buffer, so their order is the order of the calls (serial dispatch). +bool moe_sort(const int* idx, int T, int k, const int* vs, int n_vs, int* bounds, int* row_tok, int* pos) { + if (!ensure()) return false; + void* hi = strata::gpu::host_view((void*) idx); + void* hv = strata::gpu::host_view((void*) vs); + void* hb = strata::gpu::host_view((void*) bounds); + void* hr = strata::gpu::host_view((void*) row_tok); + void* hp = strata::gpu::host_view((void*) pos); + if (!hi || !hv || !hb || !hr || !hp) { std::fprintf(stderr, "dsv41 metal: moe_sort has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + const size_t t_bytes = (size_t) T * (size_t) k * 4; + mtl::Buffer bi = wrap(hi, t_bytes); + mtl::Buffer bv = wrap(hv, (size_t) n_vs * 4 ); + mtl::Buffer bb = wrap(hb, (size_t) (n_vs + 2) * 4 ); + mtl::Buffer br = wrap(hr, t_bytes); + mtl::Buffer bp = wrap(hp, t_bytes); + if (!bi.handle || !bv.handle || !bb.handle || !br.handle || !bp.handle) return false; + mtl::Pipeline pc, pf, ps; + try { + pc = g_cx->pipeline(g_lib, "ds4_moe_sort_count"); + pf = g_cx->pipeline(g_lib, "ds4_moe_sort_fill"); + ps = g_cx->pipeline(g_lib, "ds4_bounds_scan"); + } catch (...) { return false; } + // 1. counts + mtl::Dispatch d1; + d1.p = &pc; + d1.args = {mtl::Arg::B(bi, off(hi)), mtl::Arg::V((unsigned) T), mtl::Arg::V((unsigned) k), + mtl::Arg::B(bv, off(hv)), mtl::Arg::B(bb, off(hb))}; + d1.grid[0] = (uint32_t) n_vs; + d1.grid[1] = d1.grid[2] = 1; + d1.tg[0] = 256; d1.tg[1] = d1.tg[2] = 1; + g_cx->run({d1}); + // 2. the exclusive prefix over the slots + mtl::Dispatch d2; + d2.p = &ps; + d2.args = {mtl::Arg::B(bb, off(hb)), mtl::Arg::V((unsigned) n_vs)}; + d2.grid[0] = 1; + d2.grid[1] = d2.grid[2] = 1; + d2.tg[0] = 32; d2.tg[1] = d2.tg[2] = 1; + g_cx->run({d2}); + // 3. the fill + mtl::Dispatch d3; + d3.p = &pf; + d3.args = {mtl::Arg::B(bi, off(hi)), mtl::Arg::V((unsigned) T), mtl::Arg::V((unsigned) k), + mtl::Arg::B(bv, off(hv)), mtl::Arg::B(bb, off(hb)), mtl::Arg::B(br, off(hr)), mtl::Arg::B(bp, off(hp))}; + d3.grid[0] = (uint32_t) n_vs; + d3.grid[1] = d3.grid[2] = 1; + d3.tg[0] = 256; d3.tg[1] = d3.tg[2] = 1; + g_cx->run({d3}); + return true; +} + + +// ds4_kernels.cu:1690 - the Engram row decoder (fmt 2 = IQ4_NL, this model's tables) +bool engram_dequant_tok(const uint8_t* src, size_t per, int cols, int T, int hd, float* out, int fmt) { + if (T <= 0) return true; + if (!ensure()) return false; + void* hs = strata::gpu::host_view((void*) src); + void* ho = strata::gpu::host_view((void*) out); + if (!hs || !ho) { std::fprintf(stderr, "dsv41 metal: engram_dequant_tok has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + const size_t rows = (size_t) T * (size_t) cols; + mtl::Buffer bs = wrap(hs, (size_t) T * per ); + mtl::Buffer bo = wrap(ho, rows * (size_t) hd * 4 ); + if (!bs.handle || !bo.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_engram_dequant"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bs, off(hs)), mtl::Arg::V((unsigned long) per), mtl::Arg::V((unsigned) cols), + mtl::Arg::V((unsigned) hd), mtl::Arg::B(bo, off(ho)), mtl::Arg::V((int) fmt)}; + d.grid[0] = (uint32_t) rows; + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + + +// ds4_kernels.cu:341 - hc_pre_norm: a = sum_s mix[s]*h[s][e]; x = a * rsqrt(mean(a^2) + eps) (x w when given) +bool hc_pre_norm(const float* h, const float* mix, const float* w, int dim, int hc, int T, float eps, float* x) { + if (!ensure()) return false; + void* hh = strata::gpu::host_view((void*) h); + void* hm = strata::gpu::host_view((void*) mix); + void* hx = strata::gpu::host_view((void*) x); + if (!hh || !hm || !hx) { std::fprintf(stderr, "dsv41 metal: hc_pre_norm has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bh = wrap(hh, (size_t) T * hc * dim * 4 ); + mtl::Buffer bm = wrap(hm, (size_t) T * 8 * 4 ); + mtl::Buffer bx = wrap(hx, (size_t) T * dim * 4 ); + if (!bh.handle || !bm.handle || !bx.handle) return false; + mtl::Buffer bw; + size_t w_off = 0; + if (w) { + void* hw = strata::gpu::host_view((void*) w); + if (!hw) { std::fprintf(stderr, "dsv41 metal: hc_pre_norm weight has no host view\n"); return false; } + bw = wrap(hw, (size_t) dim * 4 ); + w_off = off(hw); + if (!bw.handle) return false; + } else { + bw = bx; // never read when has_w == 0 + } + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_hc_pre_norm"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bh, off(hh)), mtl::Arg::B(bm, off(hm)), mtl::Arg::B(bw, w_off), + mtl::Arg::V((unsigned) dim), mtl::Arg::V((unsigned) hc), mtl::Arg::V(eps), mtl::Arg::B(bx, off(hx)), + mtl::Arg::V((unsigned) (w ? 1 : 0))}; + d.grid[0] = (uint32_t) T; + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 1024; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + + +// ds4_gguf.cu:186 - the activation quantiser the expert matmuls consume (Q8_0: 32 values, one scale) +bool q8_quant_rows(const float* x, size_t ldx, int rows, int k, int8_t* q, float* d) { + if (rows <= 0 || k <= 0) return true; + if (!ensure()) return false; + void* hx = strata::gpu::host_view((void*) x); + void* hq = strata::gpu::host_view((void*) q); + void* hd = strata::gpu::host_view((void*) d); + if (!hx || !hq || !hd) { std::fprintf(stderr, "dsv41 metal: q8_quant has no host view\n"); return false; } + const size_t page = 16384; + // The pages the buffer itself spans: [base, base + off + bytes). A fixed slack (the +65536 the call sites used to + // carry) walks off the end of the mapping when the buffer is small and sits at an arena chunk's edge, and Metal then + // refuses the no-copy wrap ("mtl: wrap 81920 failed" - the hd scales here, a few pages). + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bx = wrap(hx, (size_t) rows * ldx * 4); + mtl::Buffer bq = wrap(hq, (size_t) rows * k); + mtl::Buffer bd = wrap(hd, (size_t) rows * (k / 32) * 4); + if (!bx.handle || !bq.handle || !bd.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_q8_quant"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bx, off(hx)), mtl::Arg::V((unsigned long) ldx), mtl::Arg::V((unsigned) k), + mtl::Arg::B(bq, off(hq)), mtl::Arg::B(bd, off(hd))}; + const uint groups = (uint32_t) ((k / 32 + 7) / 8); + dsp.grid[0] = groups; + dsp.grid[1] = (uint32_t) rows; + dsp.grid[2] = 1; + dsp.tg[0] = 256; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + return true; +} + + +// the GPU address of host weight memory (see mtl::Context::zc_map); 0: not available +uint64_t zc_map(const void* p, size_t bytes) { + if (!ensure()) return 0; + return g_cx->zc_map(p, bytes); +} + +// Several ds4_q8_quant calls in one dispatch (specs: {x, ldx, rows, k, q, d} with the engine's device pointers, which are the +// kernels' own addresses, as iq_mv's jobs carry them) +namespace { +struct HostQ8Spec { + const float* x; + uint64_t ldx; + uint32_t rows, k; + int8_t* q; + float* d; +}; +} // namespace +bool q8_quant_multi(const void* specs, int n) { + if (n <= 0) return true; + if (!ensure()) return false; + const HostQ8Spec* sp = (const HostQ8Spec*) specs; + uint32_t max_rows = 1, max_k = 32; + for (int i = 0; i < n; ++i) { + max_rows = std::max(max_rows, sp[i].rows); + max_k = std::max(max_k, sp[i].k); + } + mtl::Buffer bj = g_cx->alloc((size_t) n * sizeof(HostQ8Spec)); + std::memcpy(bj.ptr, sp, (size_t) n * sizeof(HostQ8Spec)); + static mtl::Pipeline pl; + static bool have = false; + if (!have) { + try { pl = g_cx->pipeline(g_lib, "ds4_q8_quant_multi"); } catch (...) { g_cx->free(bj); return false; } + have = true; + } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bj)}; + d.grid[0] = (max_k / 32 + 7) / 8; + d.grid[1] = max_rows; + d.grid[2] = (uint32_t) n; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + g_cx->free(bj); + return true; +} + + +// ds4_gguf.cu:371 - the MXFP4 expert matvec. The engine builds its BfJobs array on the HOST (with device pointers inside), so +// the jobs go to the kernel as an opaque constant copy: same bytes, same layout, indexed by threadgroup_position_in_grid.y. +namespace { +// mirrors ds4_gguf.cu's BfJob (and the kernels' Ds4MxJob) +struct HostBfJob { + const void* w; + uint64_t ld; + const float* x; + uint64_t xcs; + float* y; + uint64_t ycs; + const float* bias; + float alpha; + int k, n, ncol; + const int8_t* xq; + const float* xd; +}; +} // namespace + +bool mx_mv(const void* jobs_in, int nj_in, int max_n, size_t job_bytes, int fmt) { + if (nj_in <= 0 || max_n <= 0) return true; + if (!ensure()) return false; + // The matvec kernels keep 8 accumulators (the T=8 window): a job of more activations (a 16-token window) is cut into + // jobs of at most 8, each reading its own activations and writing its own rows of y. (Left whole, the columns past the + // eighth were never computed: rows of zeros in every projection of a prompt window.) + std::vector cut; + if (job_bytes == sizeof(HostBfJob)) { + const HostBfJob* src = (const HostBfJob*) jobs_in; + for (int i = 0; i < nj_in; ++i) { + const HostBfJob& q = src[i]; + const int ncol = q.ncol > 0 ? q.ncol : 1; + for (int c0 = 0; c0 < ncol; c0 += 8) { + HostBfJob b = q; + b.ncol = std::min(8, ncol - c0); + if (c0 != 0) { + if (q.x) b.x = q.x + (size_t) c0 * (size_t) q.xcs; + b.y = q.y + (size_t) c0 * (size_t) q.ycs; + if (q.xq) b.xq = q.xq + (size_t) c0 * (size_t) q.k; + if (q.xd) b.xd = q.xd + (size_t) c0 * (size_t) (q.k / 32); + } + cut.push_back(b); + } + } + } + const void* jobs = cut.empty() ? jobs_in : (const void*) cut.data(); + const int nj = cut.empty() ? nj_in : (int) cut.size(); + mtl::Buffer bj = g_cx->alloc((size_t) nj * job_bytes); + if (!bj.ptr) { std::fprintf(stderr, "dsv41 metal: mx_mv job buffer failed\n"); return false; } + std::memcpy(bj.ptr, jobs, (size_t) nj * job_bytes); + mtl::Pipeline pl; + // fmt 1 = MXFP4, 2 = Q8_0, 0 = BF16 (a dense matrix: the engram's wkv, ds4_gguf.cu:1064) + const char* fn = fmt == 2 ? "ds4_q8_mv" : (fmt == 1 ? "ds4_mx_mv" : "ds4_bf16_mv"); + try { pl = g_cx->pipeline(g_lib, fn); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bj)}; + d.grid[0] = (uint32_t) ((max_n + 7) / 8); + d.grid[1] = (uint32_t) nj; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + g_cx->free(bj); + return true; +} + + +// ds4_iq.cu:809 - the I-quant expert matvec. The jobs are the engine's own MvJob array (80 bytes each, type included), so they +// travel as an opaque constant copy and the kernel picks the decoder per job; one launch covers every type in it. +static inline void* cpu_addr(void* p); +namespace { +// mirrors ds4_iq.cuh's MvJob (and ds4_iq.metal's Ds4MvJob): 80 bytes on arm64 +struct Ds4IqMvJob { + const uint8_t* w; + uint64_t ld; + int type, k, n, ncol; + const int8_t* xq; + const float* xd; + float* y; + uint64_t ycs; + const float* bias; + float alpha; +}; +static_assert(sizeof(Ds4IqMvJob) == 80, "Ds4IqMvJob layout must match ds4_iq.cuh's MvJob"); +} // namespace + +bool iq_mv(const void* jobs, int nj) { + if (nj <= 0) return true; + if (!ensure()) return false; + const Ds4IqMvJob* src = (const Ds4IqMvJob*) jobs; + // Each job goes to the kernel of its type and column count: NC = 1, 2, 4 or 8 activations decode the row's weights once + // (ds4_*_mvn); a job of more columns is cut into pieces of 8. Other types keep the 16-column kernels. + std::map> groups; + int max_n = 0; + if (std::getenv("DS4_DBG_JOBS")) { + std::fprintf(stderr, "iq_mv call:"); + for (int i = 0; i < nj; ++i) std::fprintf(stderr, " [t%d n%d k%d c%d]", src[i].type, src[i].n, src[i].k, src[i].ncol); + std::fprintf(stderr, "\n"); + } + static const bool old_only = std::getenv("DS4_NO_MVN") != nullptr; + for (int i = 0; i < nj; ++i) { + const Ds4IqMvJob& q = src[i]; + { static const char* sk = std::getenv("DS4_SKIP_TYPES"); if (sk != nullptr && std::strstr(sk, ("," + std::to_string(q.type) + ",").c_str()) != nullptr) continue; } // (timing only) + if (q.n > max_n) max_n = q.n; + const int ncol = q.ncol > 0 ? q.ncol : 1; + const char* tn = nullptr; + switch (q.type) { + case 13: tn = "ds4_q5k_mvn"; break; + case 14: tn = "ds4_q6k_mvn"; break; // (the Unsloth sets' Q6_K dense matrices) + case 16: tn = "ds4_iq2xxs_mvn"; break; + case 17: tn = "ds4_iq2xs_mvn"; break; + case 18: tn = "ds4_iq3xxs_mvn"; break; + case 22: tn = "ds4_iq2s_mvn"; break; + case 19: tn = "ds4_iq1s_mvn"; break; // (Unsloth UD-IQ1_S / UD-IQ1_M experts: ds4_w32_iq1s) + case 29: tn = "ds4_iq1m_mvn"; break; // (the 2.00 bpw pack's IQ1_M: ds4_iq.metal's ds4_w32_iq1m) + default: break; + } + const bool grid = tn != nullptr && !old_only; + const int step = grid ? 8 : (ncol > 16 ? 16 : ncol); + for (int c0 = 0; c0 < ncol; c0 += step) { + Ds4IqMvJob b = q; + b.ncol = (c0 + step < ncol) ? step : ncol - c0; + if (c0 != 0) { + b.xq = q.xq + (size_t) c0 * (size_t) q.k; + b.xd = q.xd + (size_t) c0 * (size_t) (q.k / 32); + b.y = q.y + (size_t) c0 * (size_t) q.ycs; + } + std::string fn; + if (grid) { + const int nc = b.ncol <= 1 ? 1 : (b.ncol <= 2 ? 2 : (b.ncol <= 4 ? 4 : 8)); + fn = std::string(tn) + std::to_string(nc); + } else { + fn = (b.type == 13 || b.type == 14) ? "ds4_iq_mv_k" : "ds4_iq_mv"; + } + groups[fn].push_back(b); + } + } + if (max_n <= 0) return true; + const size_t kBatch = 64; // the CUDA host's job-array limit; one command buffer a group + for (auto& kv : groups) { + std::vector& v = kv.second; + mtl::Buffer bj = g_cx->alloc(v.size() * sizeof(Ds4IqMvJob)); + if (!bj.ptr) { + std::fprintf(stderr, "dsv41 metal: iq_mv job buffer failed\n"); + return false; + } + std::memcpy(bj.ptr, v.data(), v.size() * sizeof(Ds4IqMvJob)); + mtl::Pipeline pl; + try { + pl = g_cx->pipeline(g_lib_iq, kv.first); + } catch (...) { + g_cx->free(bj); + return false; + } + std::vector ds; + for (size_t base = 0; base < v.size(); base += kBatch) { + const size_t cnt = std::min(kBatch, v.size() - base); + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bj, base * sizeof(Ds4IqMvJob))}; + d.grid[0] = (uint32_t) ((max_n + 7) / 8); + d.grid[1] = (uint32_t) cnt; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + ds.push_back(d); + } + g_cx->run(ds); + g_cx->free(bj); + } + return true; +} + + +// ds4_iq.cu:536 gate_up_kernel / :588 down_kernel - the prompt / window path's expert GEMM (the launches at +// :958 launch_gu / :960 launch_dn). A threadgroup of 256 is 8 warps: a warp RW rows (gate / up 1, down 2), a lane the +// row's sub-blocks, the group's entries four at a pass. The group's blob and its entries are read the way the CUDA kernel +// reads them, so the plan / window pointer arrays travel as they are: every pointer here is a CPU address (unified memory) +// and a kernel may take raw gpuAddresses (as ds4_finish's f.w does). +namespace { +// (not an entry point: the generator takes only the ^bool lines) mirrors metal/src/engine/msl/ds4_iq.metal's Ds4IqJob (Metal has no 'long long': 64-bit is 'long') +struct IqJob { + const void* plan_ptr; + const void* grp_ptr; + const void* grp_start; + const void* grp_ent; + const void* n_groups; + const void* xq; + const void* xd; + const void* hq; + const void* hd; + void* h; + void* part; + int64_t tr0, tr1, tr2; + int n_embd, n_ff, type; + float limit; +}; + +static bool iq_gemm(const IqJob& j, const char* fn, int grid_x, int groups) { + if (!ensure()) return false; + mtl::Buffer bj = g_cx->alloc(sizeof(j)); + if (!bj.ptr) { + std::fprintf(stderr, "dsv41 metal: %s job buffer failed\n", fn); + return false; + } + std::memcpy(bj.ptr, &j, sizeof(j)); + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib_iq, fn); } catch (...) { g_cx->free(bj); return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bj)}; + d.grid[0] = (uint32_t) grid_x; + d.grid[1] = (uint32_t) groups; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + g_cx->free(bj); + return true; +} +} // namespace + +// gate / up of the group's entries: h = the clamped swiglu (ds4_iq.cuh's Args: tr[0] gate, tr[1] up, tr[2] down) +bool iq_gu(const unsigned long long* plan_ptr, const unsigned long long* grp_ptr, const int* grp_start, + const int* grp_ent, const int* n_groups, const size_t* tr, int n_embd, int n_ff, float limit, + const int8_t* xq, const float* xd, float* h, int type, int groups) { + if (groups <= 0) return true; + if (n_embd <= 0 || n_ff <= 0 || n_embd % 256 != 0 || n_ff % 8 != 0) return false; + if (xq == nullptr || xd == nullptr || h == nullptr) return false; + if (plan_ptr == nullptr && grp_ptr == nullptr) return false; + IqJob j{}; + j.plan_ptr = plan_ptr; + j.grp_ptr = grp_ptr; + j.grp_start = grp_start; + j.grp_ent = grp_ent; + j.n_groups = n_groups; + j.xq = xq; + j.xd = xd; + j.h = h; + j.tr0 = (int64_t) tr[0]; + j.tr1 = (int64_t) tr[1]; + j.tr2 = (int64_t) tr[2]; + j.n_embd = n_embd; + j.n_ff = n_ff; + j.type = type; + j.limit = limit; + return iq_gemm(j, "ds4_iq_gu", n_ff / 8, groups); +} + +// down of the group's entries against their int8 h: part (a row of n_embd an entry) +bool iq_dn(const unsigned long long* plan_ptr, const unsigned long long* grp_ptr, const int* grp_start, + const int* grp_ent, const int* n_groups, const size_t* tr, int n_embd, int n_ff, const int8_t* hq, + const float* hd, float* part, int type, int groups) { + if (groups <= 0) return true; + if (n_embd <= 0 || n_ff <= 0 || n_ff % 256 != 0 || n_embd % 16 != 0) return false; + if (hq == nullptr || hd == nullptr || part == nullptr) return false; + if (plan_ptr == nullptr && grp_ptr == nullptr) return false; + IqJob j{}; + j.plan_ptr = plan_ptr; + j.grp_ptr = grp_ptr; + j.grp_start = grp_start; + j.grp_ent = grp_ent; + j.n_groups = n_groups; + j.hq = hq; + j.hd = hd; + j.part = part; + j.tr0 = (int64_t) tr[0]; + j.tr1 = (int64_t) tr[1]; + j.tr2 = (int64_t) tr[2]; + j.n_embd = n_embd; + j.n_ff = n_ff; + j.type = type; + j.limit = 0.0f; + return iq_gemm(j, "ds4_iq_dn", n_embd / 16, groups); +} + + +// ds4_gguf.cu:454 mx_gate_up_kernel / :519 mx_down_kernel - the MXFP4 expert GEMM (ds4_gguf.cu's moe_gate_up / moe_down / +// win_experts). The same frame as iq_gu / iq_dn above: a threadgroup of 256 is 8 warps, a warp a row, a lane strides the +// 32-value blocks, the group's entries sixteen at a pass. The job struct is the kernel's own (metal/src/engine/msl/ +// ds4_iq.metal's Ds4MxJob); the plan / window pointer arrays travel as they are: every pointer here is a CPU address +// (unified memory) and a kernel may take raw gpuAddresses. +namespace { +// (not an entry point: the generator takes only the ^bool lines) mirrors metal/src/engine/msl/ds4_iq.metal's Ds4MxJob (Metal +// has no 'long long': 64-bit is 'long') +struct MxJob { + const void* plan_ptr; + const void* grp_ptr; + const void* grp_start; + const void* grp_ent; + const void* n_groups; + const void* xq; + const void* xd; + const void* hq; + const void* hd; + void* h; + void* part; + int64_t tr0, tr1, tr2; // ('long' there, so 8 a field and the float ends the struct on the same byte) + int n_embd, n_ff; + float limit; +}; +static_assert(sizeof(MxJob) == 128, "Ds4MxJob's layout"); + +static bool mx_gemm(const MxJob& j, const char* fn, int grid_x, int groups) { + if (!ensure()) return false; + mtl::Buffer bj = g_cx->alloc(sizeof(j)); + if (!bj.ptr) { + std::fprintf(stderr, "dsv41 metal: %s job buffer failed\n", fn); + return false; + } + std::memcpy(bj.ptr, &j, sizeof(j)); + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib_iq, fn); } catch (...) { g_cx->free(bj); return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bj)}; + d.grid[0] = (uint32_t) grid_x; + d.grid[1] = (uint32_t) groups; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + g_cx->free(bj); + return true; +} +} // namespace + +// gate / up of the group's experts: h = the clamped swiglu (ds4_gguf.cu's MxArgs: tr[0] gate, tr[1] up, tr[2] down) +bool mx_gu(const unsigned long long* plan_ptr, const unsigned long long* grp_ptr, const int* grp_start, + const int* grp_ent, const int* n_groups, const size_t* tr, int n_embd, int n_ff, float limit, + const int8_t* xq, const float* xd, float* h, int groups) { + if (groups <= 0) return true; + if (n_embd <= 0 || n_ff <= 0 || n_embd % 32 != 0 || n_ff % 8 != 0) return false; + if (xq == nullptr || xd == nullptr || h == nullptr) return false; + if (plan_ptr == nullptr && grp_ptr == nullptr) return false; + MxJob j{}; + j.plan_ptr = plan_ptr; + j.grp_ptr = grp_ptr; + j.grp_start = grp_start; + j.grp_ent = grp_ent; + j.n_groups = n_groups; + j.xq = xq; + j.xd = xd; + j.h = h; + j.tr0 = (int64_t) tr[0]; + j.tr1 = (int64_t) tr[1]; + j.tr2 = (int64_t) tr[2]; + j.n_embd = n_embd; + j.n_ff = n_ff; + j.limit = limit; + return mx_gemm(j, "ds4_mx_gu", (n_ff + 7) / 8, groups); +} + +// down of the group's experts against their int8 h: part (a row of n_embd an entry) +bool mx_dn(const unsigned long long* plan_ptr, const unsigned long long* grp_ptr, const int* grp_start, + const int* grp_ent, const int* n_groups, const size_t* tr, int n_embd, int n_ff, const int8_t* hq, + const float* hd, float* part, int groups) { + if (groups <= 0) return true; + if (n_embd <= 0 || n_ff <= 0 || n_ff % 32 != 0 || n_embd % 8 != 0) return false; + if (hq == nullptr || hd == nullptr || part == nullptr) return false; + if (plan_ptr == nullptr && grp_ptr == nullptr) return false; + MxJob j{}; + j.plan_ptr = plan_ptr; + j.grp_ptr = grp_ptr; + j.grp_start = grp_start; + j.grp_ent = grp_ent; + j.n_groups = n_groups; + j.hq = hq; + j.hd = hd; + j.part = part; + j.tr0 = (int64_t) tr[0]; + j.tr1 = (int64_t) tr[1]; + j.tr2 = (int64_t) tr[2]; + j.n_embd = n_embd; + j.n_ff = n_ff; + return mx_gemm(j, "ds4_mx_dn", (n_embd + 7) / 8, groups); +} + +// ds4_gguf.cu:577 mx_combine_kernel - a token's expert outputs summed in plan order (the experts off the device add 0), +// the CPU lane's part when its flag is set, then the shared expert's: ds4_gguf.cu:169's combine(), 128 threads a block +namespace { +// mirrors metal/src/engine/msl/ds4_init.metal's Ds4Comb (a 'long' there is int64_t here) +struct MxCombJob { + const void* plan_ptr; + const void* plan_w; + const void* part; + const void* cpu_part; + const int* cpu_flag; + const void* sh_out; + void* out; + int64_t cpu_stride, out_stride; + int k, n_embd; +}; +} // namespace + +bool mx_combine(const void* plan_ptr, const void* plan_w, const void* part, const void* cpu_part, size_t cpu_stride, + const int* cpu_flag, const void* sh_out, void* out, size_t out_stride, int T, int k, int n_embd) { + if (T <= 0 || n_embd <= 0) return true; + if (k <= 0 || plan_ptr == nullptr || plan_w == nullptr || part == nullptr || out == nullptr) return false; + if (!ensure()) return false; + MxCombJob j{}; + j.plan_ptr = plan_ptr; + j.plan_w = plan_w; + j.part = part; + j.cpu_part = cpu_part; + j.cpu_flag = cpu_flag; + j.sh_out = sh_out; + j.out = out; + j.cpu_stride = (int64_t) cpu_stride; + j.out_stride = (int64_t) out_stride; + j.k = k; + j.n_embd = n_embd; + mtl::Buffer bj = g_cx->alloc(sizeof(j)); + if (!bj.ptr) { + std::fprintf(stderr, "dsv41 metal: mx_combine job buffer failed\n"); + return false; + } + std::memcpy(bj.ptr, &j, sizeof(j)); + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_mx_combine"); } catch (...) { g_cx->free(bj); return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bj)}; + d.grid[0] = (uint32_t) ((n_embd + 127) / 128); + d.grid[1] = (uint32_t) T; + d.grid[2] = 1; + d.tg[0] = 128; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + g_cx->free(bj); + return true; +} + +// ds4_kernels.cu:1092 - the KV ring row to BF16 (E4M3 code x E8M0 group scale) +bool ring_to_ext(const void* ring, void* ext, long long p0, int n, int d, size_t row_bytes) { + if (n <= 0 || d <= 0) return true; + if (!ensure()) return false; + void* hr = strata::gpu::host_view((void*) ring); + void* he = strata::gpu::host_view(ext); + if (!hr || !he) { std::fprintf(stderr, "dsv41 metal: ring_to_ext has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer br = wrap(hr, (size_t) n * row_bytes ); + mtl::Buffer be = wrap(he, (size_t) n * d * 2 ); + if (!br.handle || !be.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_ring_to_ext"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(br, off(hr)), mtl::Arg::B(be, off(he)), mtl::Arg::V((long) p0), + mtl::Arg::V((unsigned) n), mtl::Arg::V((unsigned) d), mtl::Arg::V((unsigned long) row_bytes)}; + dsp.grid[0] = (uint32_t) n; + dsp.grid[1] = dsp.grid[2] = 1; + dsp.tg[0] = 256; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + return true; +} + + +// ds4_kernels.cu:1104 - the window back into the ring (E4M3 codes + E8M0 group scales, optional backup) +bool ext_to_ring(const void* ext, void* ring, long long p0, int n, int T, int d, size_t row_bytes, void* backup) { + if (T <= 0 || n <= 0 || d <= 0) return true; + if (!ensure()) return false; + void* he = strata::gpu::host_view((void*) ext); + void* hr = strata::gpu::host_view(ring); + if (!he || !hr) { std::fprintf(stderr, "dsv41 metal: ext_to_ring has no host view\n"); return false; } + void* hb = backup ? strata::gpu::host_view(backup) : nullptr; + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer be = wrap(he, (size_t) (n + T) * d * 2 ); + mtl::Buffer br = wrap(hr, (size_t) n * row_bytes ); + if (!be.handle || !br.handle) return false; + mtl::Buffer bb = br; // never written when backup is null + size_t b_off = 0; + if (hb) { bb = wrap(hb, (size_t) T * row_bytes ); b_off = off(hb); if (!bb.handle) return false; } + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_ext_to_ring"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(be, off(he)), mtl::Arg::B(br, off(hr)), mtl::Arg::B(bb, b_off), + mtl::Arg::V((long) p0), mtl::Arg::V((unsigned) n), mtl::Arg::V((unsigned) T), + mtl::Arg::V((unsigned) d), mtl::Arg::V((unsigned long) row_bytes), mtl::Arg::V((unsigned) (hb ? 1 : 0))}; + dsp.grid[0] = (uint32_t) T; + dsp.grid[1] = dsp.grid[2] = 1; + dsp.tg[0] = 256; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + return true; +} + + +// ds4_kernels.cu:937 - the selected compressed rows into the scratch (MXFP4 -> BF16, 16 threads a row) + +namespace { +// mirrors metal/src/engine/msl/ds4_init.metal's Ds4AttnRows (ds4_kernels.cuh:149 AttnRows, kWin = 16) +struct Ds4AttnRowsJob { + int win0[16], nwin[16], nsel[16]; + int gen, t0, win, R, kcap; + int64_t p0; +}; +static_assert(sizeof(Ds4AttnRowsJob) == 224, "Ds4AttnRowsJob layout must match the MSL struct"); +} // namespace + +// ds4_kernels.cu:954 sparse_attn_packed_kernel and :774 sparse_attn_kernel - the prompt path's sparse attention. +// The packed kernel stages rows once for G heads of a query; the BF16 kernel handles rows already in BF16 (a DSpark block). +bool sparse_attn(const float* q, const uint16_t* win, const void* ckv, const int* sel, int sel_stride, const void* rows, + const float* sink, int heads, int d, int T, float scale, float* out, int ckv_fmt) { + if (T <= 0 || heads <= 0) return true; + if (d != 512 || rows == nullptr || (ckv_fmt != 0 && ckv_fmt != 2)) { std::fprintf(stderr, "dsv41 metal: sparse_attn unsupported\n"); return false; } + Ds4AttnRowsJob ar{}; + ar.t0 = (int) T; // the two-query kernel reads the row count here on the gen == 0 path + std::memcpy(&ar, rows, sizeof(ar)); + if (q == nullptr || out == nullptr) return false; + if (ar.gen == 0 && T > 16) { std::fprintf(stderr, "dsv41 metal: sparse_attn gen0 T>16\n"); return false; } + if (!ensure()) return false; + + void* hq = cpu_addr((void*) q); + void* hw = cpu_addr((void*) win); + void* hc = cpu_addr((void*) ckv); + void* hs = cpu_addr((void*) sel); + void* hk = cpu_addr((void*) sink); + void* ho = cpu_addr((void*) out); + if (!hq || !ho) { + std::fprintf(stderr, "dsv41 metal: sparse_attn has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + + int wmax = 0; + if (ar.gen != 0) { + wmax = ar.win + T + 1; + } else { + for (int t = 0; t < T; ++t) wmax = ar.win0[t] + ar.nwin[t] > wmax ? ar.win0[t] + ar.nwin[t] : wmax; + wmax += 1; + } + if (wmax > 0 && hw == nullptr) { std::fprintf(stderr, "dsv41 metal: sparse_attn win null (wmax %d)\n", wmax); return false; } + const size_t wrows = (size_t) (wmax > 0 ? wmax : 1); + size_t crows = 0; + if (ckv_fmt == 2) { + if (ar.gen != 0 && ar.R > 0) { + crows = (size_t) ((ar.p0 + ar.t0 + T + 1) / ar.R) + 2; // every selected row is below this position + } else { + // (the rows the selection names are on the device and not read back here: an earlier build read them on the + // host while the kernel that wrote them was still queued, and sized the binding from stale zeros) + const size_t rem = ds4_gpu_alloc_remaining(hc); + crows = rem > 0 ? rem / 288 : (size_t) (ar.nsel[0] > 0 ? ar.nsel[0] : 1); + } + if (crows < 1) crows = 1; + } else { + crows = (size_t) (T > 16 ? T : 16); + } + const size_t qbytes = (size_t) T * (size_t) heads * 512 * 4; + const size_t wbytes = wrows * 512 * 2; + const size_t cbytes = crows * (ckv_fmt == 2 ? 288 : 1024); + const size_t obytes = (size_t) T * (size_t) heads * 512 * 4; + const size_t sbytes = (size_t) T * (size_t) (sel_stride > 0 ? sel_stride : 1) * 4; + mtl::Buffer bq = wrap(hq, qbytes); + mtl::Buffer bw = hw != nullptr ? wrap(hw, wbytes) : g_cx->alloc(64); + mtl::Buffer bc = hc != nullptr ? wrap(hc, cbytes) : g_cx->alloc(64); + mtl::Buffer bs = hs != nullptr ? wrap(hs, sbytes) : g_cx->alloc(4); + mtl::Buffer bk = hk != nullptr ? wrap(hk, (size_t) heads * 4) : g_cx->alloc((size_t) heads * 4); + mtl::Buffer bo = wrap(ho, obytes); + if (!bq.handle || !bw.handle || !bc.handle || !bs.handle || !bk.handle || !bo.handle) { + std::fprintf(stderr, "dsv41 metal: sparse_attn wrap failed\n"); + return false; + } + + int G = 8; + if (heads % 8 != 0) G = heads % 4 == 0 ? 4 : (heads % 2 == 0 ? 2 : 1); + while (G > 1 && (heads / G) * T < 64) G /= 2; // keep enough threadgroups on the GPU (CUDA's 2 * SMs, in effect) + + static const int attn_mode = std::getenv("DS4_ATTN_MX") ? std::atoi(std::getenv("DS4_ATTN_MX")) : 1; // 1 (default): matrix-unit kernel; 0: the scalar one; 2: both, compare + const bool use_mx = ckv_fmt == 2 && G == 8 && attn_mode != 0; + // DS4_ATTN2Q: two queries a threadgroup, sharing the staged keys (the staging is ~20 % of a 32K prompt and is per + // key, so sharing halves it). Verified bit-exact in tools/ds4/attn_harness, including the causal shape where the + // two rows have different window lengths. + static const bool attn2q_env = std::getenv("DS4_ATTN2Q") != nullptr && use_mx; + // Two queries a threadgroup need a shared window base and selection list (the kernel stages one key set for both, + // and the causal shape gives the second query a superset only when those agree). The rows carry their own + // win0/nwin/nsel, so the pairing is decided here, per dispatch, and never assumed. + bool attn2q = attn2q_env; + if (attn2q) { + for (int t = 1; t < T; ++t) { + if (ar.win0[t] != ar.win0[0] || ar.nsel[t] != ar.nsel[0] || ar.nwin[t] < ar.nwin[t - 1]) { attn2q = false; break; } + } + } + const int qpt = attn2q ? 2 : 1; + const char* fn = attn2q ? "ds4_sparse_attn_mx2" : use_mx ? "ds4_sparse_attn_mx" : ckv_fmt == 2 ? (G == 8 ? "ds4_sparse_attn_g8" : G == 4 ? "ds4_sparse_attn_g4" + : G == 2 ? "ds4_sparse_attn_g2" : "ds4_sparse_attn_g1") + : "ds4_sparse_attn_bf16"; + mtl::Pipeline pl; + try { + // DS4_ATTN_PROBE=N: the attention timing probe (see ATTN_PROBE in ds4_init.metal): 1 skips P.V, 2 skips QK. + // The numbers it produces are wrong on purpose. + const char* ap = std::getenv("DS4_ATTN_PROBE"); + if (ap != nullptr && std::atoi(ap) != 0) { + mtl::Constants cs; + cs.i(25, std::atoi(ap)); + pl = g_cx->pipeline_uncached(g_lib, fn, cs); + } else { + pl = g_cx->pipeline(g_lib, fn); + } + } catch (...) { + std::fprintf(stderr, "dsv41 metal: sparse_attn pipeline %s unavailable\n", fn); + return false; + } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bq, off(hq)), mtl::Arg::B(bw, (hw != nullptr ? off(hw) : (size_t) 0)), + mtl::Arg::B(bc, (hc != nullptr ? off(hc) : (size_t) 0)), + mtl::Arg::B(bs, (hs != nullptr ? off(hs) : (size_t) 0)), + mtl::Arg::V((unsigned) (sel_stride > 0 ? sel_stride : 1)), mtl::Arg::V(ar), + mtl::Arg::B(bk, (hk != nullptr ? off(hk) : (size_t) 0)), mtl::Arg::V((unsigned) heads), mtl::Arg::V(scale), + mtl::Arg::B(bo, off(ho))}; + dsp.grid[0] = (uint32_t) (ckv_fmt == 2 ? heads / G : heads); + dsp.grid[1] = (uint32_t) ((T + qpt - 1) / qpt); + dsp.grid[2] = 1; + dsp.tg[0] = 256; + dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + // the stand-ins for absent inputs are fresh Metal buffers (16 KB of address space each): give them back + if (hw == nullptr) g_cx->free(bw); + if (hc == nullptr) g_cx->free(bc); + if (hs == nullptr) g_cx->free(bs); + if (hk == nullptr) g_cx->free(bk); + if (use_mx && attn_mode == 2) { // debugging: the scalar kernel on the same inputs, then the difference + mtl::Pipeline ps; + try { ps = g_cx->pipeline(g_lib, "ds4_sparse_attn_g8"); } catch (...) { return true; } + mtl::Buffer bo2 = g_cx->alloc(obytes); + mtl::Dispatch d2 = dsp; + d2.p = &ps; + d2.args.back() = mtl::Arg::B(bo2, 0); + g_cx->run({d2}); + g_cx->flush_runs(); + const float* a = (const float*) ho; + const float* c = (const float*) bo2.ptr; + double mx_abs = 0, mx_ref = 0, sum_sq = 0, sum_ref = 0; + for (size_t i = 0; i < obytes / 4; ++i) { + const double df = std::fabs((double) a[i] - (double) c[i]); + mx_abs = std::max(mx_abs, df); + mx_ref = std::max(mx_ref, (double) std::fabs(c[i])); + sum_sq += df * df; + sum_ref += (double) c[i] * (double) c[i]; + } + static double worst_rel = 0; static long calls = 0; + const double rel = std::sqrt(sum_sq / (sum_ref + 1e-30)); + worst_rel = std::max(worst_rel, rel); + if (++calls % 50 == 1) std::fprintf(stderr, "attn mx vs scalar: call %ld T %d max abs %.3g (max |ref| %.3g) rms rel %.3g worst rms rel %.3g\n", calls, T, mx_abs, mx_ref, rel, worst_rel); + g_cx->free(bo2); + } + return true; +} + + +// ds4_kernels.cu:461 route_kernel - sqrt(softplus(logits)) top k with bias, then the unbiased weights. vl is a VlBits +// (1024 bytes) or null; bias_vl applies to the rows its bits select. +bool route(const float* logits, const float* bias, int n, int k, int T, float route_scale, int* idx, float* w, + const float* bias_vl, const void* vl) { + if (T <= 0) return true; + if (logits == nullptr || bias == nullptr || idx == nullptr || w == nullptr || n <= 0 || n > 1024 || k <= 0 || k > 8) + return false; + if (!ensure()) return false; + void* hl = cpu_addr((void*) logits); + void* hb = cpu_addr((void*) bias); + void* hi = cpu_addr((void*) idx); + void* hw = cpu_addr((void*) w); + void* hv = bias_vl != nullptr ? cpu_addr((void*) bias_vl) : nullptr; + if (!hl || !hb || !hi || !hw) { + std::fprintf(stderr, "dsv41 metal: route has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bl = wrap(hl, (size_t) T * 1024 * 4); + mtl::Buffer bb = wrap(hb, (size_t) n * 4); + mtl::Buffer bi = wrap(hi, (size_t) T * 8 * 4); + mtl::Buffer bw = wrap(hw, (size_t) T * 8 * 4); + mtl::Buffer bv = hv != nullptr ? wrap(hv, (size_t) n * 4) : g_cx->alloc(4); + mtl::Buffer bbv = g_cx->alloc(1024); + if (!bl.handle || !bb.handle || !bi.handle || !bw.handle || !bv.handle || !bbv.ptr) return false; + const bool has_vl = vl != nullptr; + if (has_vl) std::memcpy(bbv.ptr, vl, 1024); + mtl::Pipeline pl; + try { + pl = g_cx->pipeline(g_lib, "ds4_route"); + } catch (...) { + std::fprintf(stderr, "dsv41 metal: route pipeline unavailable\n"); + return false; + } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bl, off(hl)), mtl::Arg::B(bb, off(hb)), mtl::Arg::V((unsigned) n), mtl::Arg::V((unsigned) k), + mtl::Arg::V(route_scale), mtl::Arg::B(bi, off(hi)), mtl::Arg::B(bw, off(hw)), + mtl::Arg::B(bv, hv != nullptr ? off(hv) : 0), mtl::Arg::V(has_vl ? 1u : 0u), mtl::Arg::B(bbv)}; + dsp.grid[0] = (uint32_t) T; + dsp.grid[1] = dsp.grid[2] = 1; + dsp.tg[0] = 512; + dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + g_cx->free(bbv); + if (hv == nullptr) g_cx->free(bv); + return true; +} + + +namespace { +struct Ds4GatherJob { + int src[16 * 8]; +}; +struct Ds4CombineJob { + int n[16]; + int row[16][8]; + float w[16][8]; + int cpu[16]; +}; +static_assert(sizeof(Ds4GatherJob) == 512, "Ds4GatherJob layout must match the MSL struct"); +static_assert(sizeof(Ds4CombineJob) == 1152, "Ds4CombineJob layout must match the MSL struct"); +} // namespace + +// ds4_kernels.cu:1336 gather_kernel - out[r] = x[g.src[r]] for n rows of cols floats +bool gather_rows(const float* x, const void* g, int n, int cols, float* out) { + if (n <= 0 || cols <= 0) return true; + if (x == nullptr || g == nullptr || out == nullptr || n > 16 * 8) return false; + if (!ensure()) return false; + void* hx = cpu_addr((void*) x); + void* ho = cpu_addr((void*) out); + if (!hx || !ho) return false; + Ds4GatherJob gj{}; + std::memcpy(&gj, g, sizeof(gj)); + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + int max_src = 0; + for (int i = 0; i < n; ++i) max_src = gj.src[i] > max_src ? gj.src[i] : max_src; + mtl::Buffer bx = wrap(hx, (size_t) (max_src + 1) * cols * 4); // src rows may be beyond the n output rows + mtl::Buffer bo = wrap(ho, (size_t) n * cols * 4); + if (!bx.handle || !bo.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_gather"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bx, off(hx)), mtl::Arg::V(gj), mtl::Arg::V((unsigned) cols), mtl::Arg::B(bo, off(ho))}; + dsp.grid[0] = (uint32_t) ((cols + 255) / 256); + dsp.grid[1] = (uint32_t) n; + dsp.grid[2] = 1; + dsp.tg[0] = 256; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + return true; +} + +// ds4_kernels.cu:1342 combine_kernel - out[t] = sum_j w[t][j] ed[row[t][j]] + cpu[t] (when set) + sh[t] +bool combine(const float* ed, const float* cpu, const float* sh, const void* c, int T, int dim, float* out) { + if (T <= 0 || dim <= 0) return true; + if (ed == nullptr || c == nullptr || out == nullptr || T > 16) return false; + if (!ensure()) return false; + void* he = cpu_addr((void*) ed); + void* hc = cpu != nullptr ? cpu_addr((void*) cpu) : nullptr; + void* hs = sh != nullptr ? cpu_addr((void*) sh) : nullptr; + void* ho = cpu_addr((void*) out); + if (!he || !ho) return false; + Ds4CombineJob cj{}; + std::memcpy(&cj, c, sizeof(cj)); + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + int max_row = 0; + for (int t = 0; t < T; ++t) + for (int j = 0; j < cj.n[t]; ++j) max_row = cj.row[t][j] > max_row ? cj.row[t][j] : max_row; + mtl::Buffer be = wrap(he, (size_t) (max_row + 1) * dim * 4); // rows may be beyond T * kPicks + mtl::Buffer bc = hc != nullptr ? wrap(hc, (size_t) T * dim * 4) : g_cx->alloc(4); + mtl::Buffer bs = hs != nullptr ? wrap(hs, (size_t) T * dim * 4) : g_cx->alloc(4); + mtl::Buffer bo = wrap(ho, (size_t) T * dim * 4); + if (!be.handle || !bc.handle || !bs.handle || !bo.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_combine"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(be, off(he)), mtl::Arg::B(bc, hc != nullptr ? off(hc) : 0), + mtl::Arg::B(bs, hs != nullptr ? off(hs) : 0), mtl::Arg::V(cj), mtl::Arg::V((unsigned) dim), + mtl::Arg::V(hc != nullptr ? 1u : 0u), mtl::Arg::V(hs != nullptr ? 1u : 0u), mtl::Arg::B(bo, off(ho))}; + dsp.grid[0] = (uint32_t) ((dim + 255) / 256); + dsp.grid[1] = (uint32_t) T; + dsp.grid[2] = 1; + dsp.tg[0] = 256; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + if (hc == nullptr) g_cx->free(bc); + if (hs == nullptr) g_cx->free(bs); + return true; +} + + +// On Apple silicon every device pointer is a CPU address (unified memory), so a pointer the gpu allocator does not know about +// (the engine builds some of its small arrays itself) still wraps correctly: host_view when it has one, the pointer otherwise. +static inline void* cpu_addr(void* p) { + void* h = strata::gpu::host_view(p); + return h ? h : p; +} + +// ds4_kernels.cu:1286 engram_dequant_kernel - the Engram rows (fmt 2 = IQ4_NL, the tables this model's pack carries; the +// other formats are left to the CUDA launch, which reports the line) +bool engram_dequant(const uint8_t* rows, const uint8_t* scales, int n, int hd, float* out, int fmt) { + (void) scales; // (fmt 2 keeps every block's scale in its own 18 bytes) + if (n <= 0 || hd <= 0) return true; + if (fmt != 2 || rows == nullptr || out == nullptr) return false; + if (!ensure()) return false; + void* hr = cpu_addr((void*) rows); + void* ho = cpu_addr((void*) out); + if (!hr || !ho) { std::fprintf(stderr, "dsv41 metal: engram_dequant has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + // (the kernel strides rows by hd bytes, as the CUDA kernel does, though an IQ4_NL row only fills hd/32*18 of them: + // wrapping n * 18/32 * hd left the last rows outside the buffer - garbage scales, measured as inf at row 257 of 288) + mtl::Buffer br = wrap(hr, (size_t) n * (size_t) hd); + mtl::Buffer bo = wrap(ho, (size_t) n * (size_t) hd * 4); + if (!br.handle || !bo.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_engram_dequant_rows"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(br, off(hr)), mtl::Arg::V((unsigned) hd), mtl::Arg::B(bo, off(ho)), mtl::Arg::V(fmt)}; + d.grid[0] = (uint32_t) n; + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + +// ds4_kernels.cu:1311 engram_gate_kernel - the Engram gate over a (token, copy) row: h += sigmoid(normalised dot) * val +bool engram_gate(float* h, const float* kv, const float* w, int hc, int dim, int T, float eps) { + if (T <= 0 || hc <= 0 || dim <= 0) return true; + if (h == nullptr || kv == nullptr || w == nullptr) return false; + if (!ensure()) return false; + void* hh = cpu_addr((void*) h); + void* hk = cpu_addr((void*) kv); + void* hw = cpu_addr((void*) w); + if (!hh || !hk || !hw) { std::fprintf(stderr, "dsv41 metal: engram_gate has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bh = wrap(hh, (size_t) T * (size_t) hc * (size_t) dim * 4); + mtl::Buffer bk = wrap(hk, (size_t) T * (size_t) (hc + 1) * (size_t) dim * 4); + mtl::Buffer bw = wrap(hw, (size_t) hc * (size_t) dim * 4); + if (!bh.handle || !bk.handle || !bw.handle) { + std::fprintf(stderr, "dsv41 metal: engram_gate wrap failed (h %p kv %p w %p; %p %p %p)\n", hh, hk, hw, + bh.handle, bk.handle, bw.handle); + return false; + } + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_engram_gate"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bh, off(hh)), mtl::Arg::B(bk, off(hk)), mtl::Arg::B(bw, off(hw)), + mtl::Arg::V((unsigned) hc), mtl::Arg::V((unsigned) dim), mtl::Arg::V(eps)}; + d.grid[0] = (uint32_t) hc; + d.grid[1] = (uint32_t) T; + d.grid[2] = 1; + d.tg[0] = 1024; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + +namespace { +// mirrors metal/src/engine/msl/ds4_init.metal's Ds4Rows (ds4_kernels.cuh:115 Rows, kWin = 16) +struct Ds4RowsJob { + int n[16]; + int k[16]; + int gen, R, kcap, ccap, blk; + int64_t p0; +}; +} // namespace + +// ds4_kernels.cu:524 compress_kernel - the window compressor (ratio > 1): each token's kv / score row into its slot, and +// every ratio'th token pooled by a per-value softmax over the slots, scaled by the row RMS and the norm weights +bool compress_win(const float* kv, const float* score, float* kv_state, float* score_state, int ratio, int p0, int T, + int d, const float* w, float eps, float* latent) { + if (T <= 0 || d <= 0 || ratio <= 0) return true; + if (d > 2 * 512) return false; // (the kernel holds d <= 2 * blockDim) + if (kv == nullptr || score == nullptr || kv_state == nullptr || score_state == nullptr || w == nullptr || + latent == nullptr) + return false; + if (!ensure()) return false; + void* hkv = cpu_addr((void*) kv); + void* hsc = cpu_addr((void*) score); + void* hks = cpu_addr((void*) kv_state); + void* hss = cpu_addr((void*) score_state); + void* hw = cpu_addr((void*) w); + void* hla = cpu_addr((void*) latent); + if (!hkv || !hsc || !hks || !hss || !hw || !hla) { + std::fprintf(stderr, "dsv41 metal: compress_win has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bkv = wrap(hkv, (size_t) T * (size_t) d * 4); + mtl::Buffer bsc = wrap(hsc, (size_t) T * (size_t) d * 4); + mtl::Buffer bks = wrap(hks, (size_t) ratio * (size_t) d * 4); + mtl::Buffer bss = wrap(hss, (size_t) ratio * (size_t) d * 4); + mtl::Buffer bw = wrap(hw, (size_t) d * 4); + mtl::Buffer bla = wrap(hla, (size_t) T * (size_t) d * 4); + if (!bkv.handle || !bsc.handle || !bks.handle || !bss.handle || !bw.handle || !bla.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_compress"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bkv, off(hkv)), mtl::Arg::B(bsc, off(hsc)), mtl::Arg::B(bks, off(hks)), + mtl::Arg::B(bss, off(hss)), mtl::Arg::V(ratio), mtl::Arg::V(p0), mtl::Arg::V(T), + mtl::Arg::V((unsigned) d), mtl::Arg::B(bw, off(hw)), mtl::Arg::V(eps), mtl::Arg::B(bla, off(hla))}; + dsp.grid[0] = 1; + dsp.grid[1] = dsp.grid[2] = 1; + dsp.tg[0] = 512; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + return true; +} + +// ds4_kernels.cu:582 index_score_kernel<128, kfmt> - the indexer's scores: a position's key row against every head's +// query, relu, the heads' weighted sum (rows = ds4_kernels.cuh's Rows, by value there and copied here) +bool index_score(const float* q, const float* wq, const void* kc, int heads, int hd, int T, int n, const int* cand, + int cand_stride, int blk, const void* rows, float* out, long long out_stride, int kfmt) { + if (n <= 0 || T <= 0) return true; + if (hd != 128 || heads <= 0 || heads > 32 || (kfmt != 3 && kfmt != 0)) return false; + if (q == nullptr || wq == nullptr || kc == nullptr || out == nullptr || rows == nullptr) return false; + if (!ensure()) return false; + void* hq = cpu_addr((void*) q); + void* hwq = cpu_addr((void*) wq); + void* hkc = cpu_addr((void*) kc); + void* hout = cpu_addr((void*) out); + if (!hq || !hwq || !hkc || !hout) { + std::fprintf(stderr, "dsv41 metal: index_score has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + const size_t kb = kfmt == 3 ? (size_t) ((hd / 2 + hd / 32 + 15) / 16 * 16) : (size_t) hd * 2; + mtl::Buffer bq = wrap(hq, (size_t) T * (size_t) heads * (size_t) hd * 4); + mtl::Buffer bwq = wrap(hwq, (size_t) T * 64 * 4); + mtl::Buffer bkc = wrap(hkc, (size_t) n * kb); + mtl::Buffer bo = wrap(hout, ((size_t) (T - 1) * (size_t) out_stride + (size_t) n) * 4); + if (!bq.handle || !bwq.handle || !bkc.handle || !bo.handle) return false; + const size_t cbytes = (size_t) T * (size_t) (cand_stride > 0 ? cand_stride : 1) * 4; + mtl::Buffer bc = cand != nullptr ? wrap(cpu_addr((void*) cand), cbytes) : g_cx->alloc(4); + if (!bc.handle) return false; + // Rows is a host struct: it travels as a copy (no gpu address) + mtl::Buffer br = g_cx->alloc(sizeof(Ds4RowsJob)); + if (!br.ptr) { std::fprintf(stderr, "dsv41 metal: index_score rows buffer failed\n"); return false; } + std::memcpy(br.ptr, rows, sizeof(Ds4RowsJob)); + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_index_score"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bq, off(hq)), mtl::Arg::B(bwq, off(hwq)), mtl::Arg::B(bkc, off(hkc)), + mtl::Arg::B(bc, cand != nullptr ? off(cpu_addr((void*) cand)) : 0), mtl::Arg::B(bo, off(hout)), + mtl::Arg::B(br), mtl::Arg::V(heads), mtl::Arg::V(n), mtl::Arg::V(cand_stride), mtl::Arg::V(blk), + mtl::Arg::V((long) out_stride), mtl::Arg::V(kfmt), mtl::Arg::V(cand != nullptr ? 1 : 0)}; + // DS4_IDX_MMA: the matrix-unit indexer (ds4_index_score_mma). It stages 16 fp4 keys a threadgroup and runs the + // 32 x 128 x 16 head-dot product through simdgroup_matrix - 3.17-6.10x the scalar loop in the standalone harness, + // with every output matching to half precision. Opt-in until the gate and KL say the rounding is acceptable. + // On by default: measured +7.7 to +12 % on the 32K prompt and +2 % at 3K, with the KL unchanged on both packs + // (0.11018 and 0.17087). DS4_NO_IDX_MMA=1 restores the scalar indexer; on the 2.00 pack the score hash moves + // (e81282a3 -> 0800179d) because half rounding breaks token ties, while KL and top-1 are identical. + static const bool idx_mma = std::getenv("DS4_NO_IDX_MMA") == nullptr; + if (idx_mma && cand == nullptr && kfmt == 3) { + mtl::Pipeline pm; + try { pm = g_cx->pipeline(g_lib, "ds4_index_score_mma"); } catch (...) { pm = mtl::Pipeline{}; } + if (pm.handle) { + mtl::Dispatch dm; + dm.p = ± + dm.args = {mtl::Arg::B(bq, off(hq)), mtl::Arg::B(bwq, off(hwq)), mtl::Arg::B(bkc, off(hkc)), + mtl::Arg::B(bc, 0), mtl::Arg::B(bo, off(hout)), mtl::Arg::B(br), + mtl::Arg::V(heads), mtl::Arg::V(n), mtl::Arg::V((int) T), mtl::Arg::V((long) out_stride)}; + dm.grid[0] = (uint32_t) ((n + 15) / 16); + dm.grid[1] = (uint32_t) T; + dm.grid[2] = 1; + dm.tg[0] = 256; dm.tg[1] = dm.tg[2] = 1; + g_cx->run({dm}); + g_cx->free(br); + if (bc.handle && cand != nullptr) g_cx->free(bc); + return true; + } + } + dsp.grid[0] = (uint32_t) ((n + 255) / 256); + dsp.grid[1] = (uint32_t) T; + dsp.grid[2] = 1; + dsp.tg[0] = 256; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + g_cx->free(br); + if (cand == nullptr) g_cx->free(bc); + return true; +} + +// ds4_kernels.cu:401 gemv_f16_rows_kernel - F16 weights against F32 activations, a warp a row of the matrix and 2 tokens +// a block (the indexer's wproj: ds4_model.cu:7743) +bool gemv_f16_rows(const uint16_t* W, const float* x, size_t ldx, int n, int k, int T, float alpha, float* y, size_t ldy) { + if (T <= 0 || n <= 0 || k <= 0) return true; + if (k % 8 != 0 || W == nullptr || x == nullptr || y == nullptr) return false; + if (!ensure()) return false; + void* hw = cpu_addr((void*) W); + void* hx = cpu_addr((void*) x); + void* hy = cpu_addr((void*) y); + if (!hw || !hx || !hy) { + std::fprintf(stderr, "dsv41 metal: gemv_f16_rows has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bw = wrap(hw, (size_t) n * (size_t) k * 2); + mtl::Buffer bx = wrap(hx, ((size_t) (T - 1) * ldx + (size_t) k) * 4); + mtl::Buffer by = wrap(hy, ((size_t) (T - 1) * ldy + (size_t) n) * 4); + if (!bw.handle || !bx.handle || !by.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_gemv_f16_rows"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bw, off(hw)), mtl::Arg::B(bx, off(hx)), mtl::Arg::V((unsigned long) ldx), mtl::Arg::V(n), + mtl::Arg::V(k), mtl::Arg::V(T), mtl::Arg::V(alpha), mtl::Arg::B(by, off(hy)), + mtl::Arg::V((unsigned long) ldy)}; + dsp.grid[0] = (uint32_t) ((n + 7) / 8); + dsp.grid[1] = (uint32_t) ((T + 1) / 2); + dsp.grid[2] = 1; + dsp.tg[0] = 256; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + return true; +} + +// ds4_kernels.cu:650 topk_kernel - a token's k largest indexer scores (four 8-bit radix passes, then the picks in index +// order); Rows by value there and copied here +bool topk(const float* vals, long long stride, int T, const void* rows, const int* cand, int cand_stride, int blk, int* out, + int out_stride) { + if (T <= 0) return true; + if (vals == nullptr || out == nullptr || rows == nullptr) return false; + if (!ensure()) return false; + void* hv = cpu_addr((void*) vals); + void* ho = cpu_addr((void*) out); + if (!hv || !ho) { + std::fprintf(stderr, "dsv41 metal: topk has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bv = wrap(hv, ((size_t) (T - 1) * (size_t) stride + 1) * 4); + mtl::Buffer bo = wrap(ho, ((size_t) (T - 1) * (size_t) out_stride + 1) * 4); + if (!bv.handle || !bo.handle) return false; + const size_t cbytes = (size_t) T * (size_t) (cand_stride > 0 ? cand_stride : 1) * 4; + mtl::Buffer bc = cand != nullptr ? wrap(cpu_addr((void*) cand), cbytes) : g_cx->alloc(4); + if (!bc.handle) return false; + mtl::Buffer br = g_cx->alloc(sizeof(Ds4RowsJob)); + if (!br.ptr) { std::fprintf(stderr, "dsv41 metal: topk rows buffer failed\n"); return false; } + std::memcpy(br.ptr, rows, sizeof(Ds4RowsJob)); + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_topk"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bv, off(hv)), mtl::Arg::V((long) stride), mtl::Arg::B(br), + mtl::Arg::B(bc, cand != nullptr ? off(cpu_addr((void*) cand)) : 0), mtl::Arg::V(cand_stride), + mtl::Arg::V(blk), mtl::Arg::B(bo, off(ho)), mtl::Arg::V(out_stride), + mtl::Arg::V(cand != nullptr ? 1 : 0)}; + dsp.grid[0] = (uint32_t) T; + dsp.grid[1] = dsp.grid[2] = 1; + dsp.tg[0] = 1024; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + g_cx->free(br); + if (cand == nullptr) g_cx->free(bc); + return true; +} + +// ds4_kernels.cu:751 block_max_kernel - a block of the indexer's positions and its largest score (the newest pinned at +inf) +bool block_max(const float* score, long long stride, int T, const void* rows, int blk, float* bm, long long bm_stride, + int nb_max) { + if (T <= 0 || nb_max <= 0 || blk <= 0) return true; + if (score == nullptr || bm == nullptr || rows == nullptr) return false; + if (!ensure()) return false; + void* hs = cpu_addr((void*) score); + void* hb = cpu_addr((void*) bm); + if (!hs || !hb) { + std::fprintf(stderr, "dsv41 metal: block_max has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bs = wrap(hs, ((size_t) (T - 1) * (size_t) stride + (size_t) nb_max * (size_t) blk) * 4); + mtl::Buffer bb = wrap(hb, ((size_t) (T - 1) * (size_t) bm_stride + (size_t) nb_max) * 4); + if (!bs.handle || !bb.handle) return false; + mtl::Buffer br = g_cx->alloc(sizeof(Ds4RowsJob)); + if (!br.ptr) { std::fprintf(stderr, "dsv41 metal: block_max rows buffer failed\n"); return false; } + std::memcpy(br.ptr, rows, sizeof(Ds4RowsJob)); + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_block_max"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bs, off(hs)), mtl::Arg::V((long) stride), mtl::Arg::B(br), mtl::Arg::V(blk), + mtl::Arg::B(bb, off(hb)), mtl::Arg::V((long) bm_stride)}; + dsp.grid[0] = (uint32_t) ((nb_max + 255) / 256); + dsp.grid[1] = (uint32_t) T; + dsp.grid[2] = 1; + dsp.tg[0] = 256; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + g_cx->free(br); + return true; +} + +// ds4_kernels.cu:1141 hc_mean_kernel - a token's mean over its hc streams (grid (ceil(dim / 256), T), 256 threads) +bool hc_mean(const float* h, int T, int hc, int dim, float* out, long long out_stride) { + if (T <= 0 || hc <= 0 || dim <= 0) return true; + if (h == nullptr || out == nullptr) return false; + if (!ensure()) return false; + void* hh = cpu_addr((void*) h); + void* ho = cpu_addr((void*) out); + if (!hh || !ho) { + std::fprintf(stderr, "dsv41 metal: hc_mean has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bh = wrap(hh, (size_t) T * (size_t) hc * (size_t) dim * 4); + mtl::Buffer bo = wrap(ho, ((size_t) (T - 1) * (size_t) out_stride + (size_t) dim) * 4); + if (!bh.handle || !bo.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_hc_mean"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bh, off(hh)), mtl::Arg::V(hc), mtl::Arg::V(dim), mtl::Arg::B(bo, off(ho)), + mtl::Arg::V((long) out_stride)}; + dsp.grid[0] = (uint32_t) ((dim + 255) / 256); + dsp.grid[1] = (uint32_t) T; + dsp.grid[2] = 1; + dsp.tg[0] = 256; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + return true; +} + +// ds4_kernels.cu:1150 iota_rows_kernel - out[t][i] = i +// ds4_kernels.cu:1156 markov_kernel / :1265 confidence_kernel - the drafter's head. markov_add is a warp a vocabulary +// row (the previous token's E row staged in the threadgroup, then E . H[v] added to logits[v]); confidence a threadgroup +// a token. Both jobs are the kernels' own structs (metal/src/engine/msl/ds4_init.metal's Ds4MarkovJob / Ds4ConfJob). +namespace { +// (not an entry point: the generator takes only the ^bool lines) mirrors ds4_init.metal's Ds4MarkovJob / Ds4ConfJob +struct MarkovJob { + void* logits; + const void* E; + const void* H; + const void* tok; + int V, R; +}; +static_assert(sizeof(MarkovJob) == 40, "Ds4MarkovJob's layout"); + +struct ConfJob { + const void* x; + const void* E; + const void* tok; + const void* w; + void* out; + int dim, R; +}; +static_assert(sizeof(ConfJob) == 48, "Ds4ConfJob's layout"); +} // namespace + +bool markov_add(const float* logits, const uint16_t* E, const uint16_t* H, const int* tok, int V, int R) { + if (logits == nullptr || E == nullptr || H == nullptr || tok == nullptr) return false; + if (V <= 0 || R <= 0 || R > 1024) return false; // (the kernel's staged E row holds 1024 floats) + if (!ensure()) return false; + MarkovJob j{}; + j.logits = const_cast(logits); + j.E = E; + j.H = H; + j.tok = tok; + j.V = V; + j.R = R; + mtl::Buffer bj = g_cx->alloc(sizeof(j)); + if (!bj.ptr) { + std::fprintf(stderr, "dsv41 metal: markov_add job buffer failed\n"); + return false; + } + std::memcpy(bj.ptr, &j, sizeof(j)); + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_markov_add"); } catch (...) { g_cx->free(bj); return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bj)}; + d.grid[0] = (uint32_t) ((V + 7) / 8); + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + g_cx->free(bj); + return true; +} + +bool confidence(const float* x, const uint16_t* E, const int* tok, const float* w, int T, int dim, int R, float* out) { + if (x == nullptr || E == nullptr || tok == nullptr || w == nullptr || out == nullptr) return false; + if (T <= 0 || dim <= 0 || R <= 0) return false; + if (!ensure()) return false; + ConfJob j{}; + j.x = x; + j.E = E; + j.tok = tok; + j.w = w; + j.out = out; + j.dim = dim; + j.R = R; + mtl::Buffer bj = g_cx->alloc(sizeof(j)); + if (!bj.ptr) { + std::fprintf(stderr, "dsv41 metal: confidence job buffer failed\n"); + return false; + } + std::memcpy(bj.ptr, &j, sizeof(j)); + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_confidence"); } catch (...) { g_cx->free(bj); return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bj)}; + d.grid[0] = (uint32_t) T; + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 1024; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + g_cx->free(bj); + return true; +} + +bool iota_rows(int* out, int T, int n, int stride) { + if (T <= 0 || n <= 0) return true; + if (out == nullptr) return false; + if (!ensure()) return false; + void* ho = cpu_addr((void*) out); + if (!ho) { + std::fprintf(stderr, "dsv41 metal: iota_rows has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bo = wrap(ho, ((size_t) (T - 1) * (size_t) stride + (size_t) n) * 4); + if (!bo.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_iota_rows"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bo, off(ho)), mtl::Arg::V(n), mtl::Arg::V(stride)}; + dsp.grid[0] = (uint32_t) T; + dsp.grid[1] = dsp.grid[2] = 1; + dsp.tg[0] = 64; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + return true; +} + +// vendor/sampler.cu:93 sampler_greedy_kernel - the plain argmax over the vocabulary, one block (1024) a token, the +// repetition penalties from the history window's bitmap (the caller takes this path when greedy or temperature <= 0) +bool sample_greedy(const float* logits, int n_tokens, int n_vocab, const int* history, int history_len, int penalty_last_n, + float penalty_repeat, float penalty_freq, float penalty_present, const float* logit_bias, int* out) { + if (n_tokens <= 0 || n_vocab <= 0) return true; + if (logits == nullptr || out == nullptr) return false; + if ((n_vocab + 31) / 32 > 4096) return false; // (the kernel's threadgroup bitmap) + if (!ensure()) return false; + void* hl = cpu_addr((void*) logits); + void* ho = cpu_addr((void*) out); + if (!hl || !ho) { + std::fprintf(stderr, "dsv41 metal: sample_greedy has no host view\n"); + return false; + } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bl = wrap(hl, (size_t) n_tokens * (size_t) n_vocab * 4); + mtl::Buffer bo = wrap(ho, (size_t) n_tokens * 4); + if (!bl.handle || !bo.handle) return false; + const size_t hbytes = (size_t) n_tokens * (size_t) (history_len > 0 ? history_len : 1) * 4; + mtl::Buffer bh = history != nullptr ? wrap(cpu_addr((void*) history), hbytes) : g_cx->alloc(4); + mtl::Buffer bb = logit_bias != nullptr ? wrap(cpu_addr((void*) logit_bias), (size_t) n_vocab * 4) : g_cx->alloc(4); + if (!bh.handle || !bb.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_sampler_greedy"); } catch (...) { return false; } + mtl::Dispatch dsp; + dsp.p = &pl; + dsp.args = {mtl::Arg::B(bl, off(hl)), mtl::Arg::V(n_vocab), + mtl::Arg::B(bh, history != nullptr ? off(cpu_addr((void*) history)) : 0), + mtl::Arg::V(history_len), mtl::Arg::V(penalty_last_n), mtl::Arg::V(penalty_repeat), + mtl::Arg::V(penalty_freq), mtl::Arg::V(penalty_present), + mtl::Arg::B(bb, logit_bias != nullptr ? off(cpu_addr((void*) logit_bias)) : 0), + mtl::Arg::V(logit_bias != nullptr ? 1 : 0), mtl::Arg::B(bo, off(ho))}; + dsp.grid[0] = (uint32_t) n_tokens; + dsp.grid[1] = dsp.grid[2] = 1; + dsp.tg[0] = 1024; dsp.tg[1] = dsp.tg[2] = 1; + g_cx->run({dsp}); + if (history == nullptr) g_cx->free(bh); + if (logit_bias == nullptr) g_cx->free(bb); + return true; +} + +bool unpack_sel(const void* ckv, const int* sel, int sel_stride, int gen, long long p0, int t0, int R, int kcap, + const int* nsel, int cap, void* out, int T, int nsel_count, size_t row_bytes) { + if (T <= 0 || sel_stride <= 0 || nsel_count <= 0) return true; + // nothing selected yet (a prompt at position 0): the scratch stays untouched, which is what the CUDA launch would do with + // no rows to write + if (ckv == nullptr || sel == nullptr || out == nullptr) return true; + if (!ensure()) return false; + void* hc = cpu_addr((void*) ckv); + void* hs = cpu_addr((void*) sel); + void* ho = cpu_addr((void*) out); + if (!hc || !hs || !ho) { std::fprintf(stderr, "dsv41 metal: unpack_sel has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + int maxrow = -1; + if (gen != 0 && R > 0) { + maxrow = (int) ((p0 + t0 + T) / R) + 1; + } else if (gen == 0 && nsel != nullptr) { + for (int t = 0; t < T && t < nsel_count; ++t) + for (int i = 0; i < nsel[t]; ++i) { + const int r = sel[(size_t) t * (size_t) sel_stride + (size_t) i]; + if (r > maxrow) maxrow = r; + } + } + mtl::Buffer bc = wrap(hc, (size_t) (maxrow + 1 > 0 ? maxrow + 1 : 1) * row_bytes); + mtl::Buffer bsl = wrap(hs, (size_t) T * sel_stride * 4 ); + mtl::Buffer bo = wrap(ho, (size_t) T * cap * 512 * 2 ); + if (!bc.handle || !bsl.handle || !bo.handle) return false; + mtl::Buffer bn = g_cx->alloc((size_t) nsel_count * 4); + if (!bn.ptr) return false; + std::memcpy(bn.ptr, nsel, (size_t) nsel_count * 4); + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_unpack_sel"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bc, off(hc)), mtl::Arg::B(bsl, off(hs)), mtl::Arg::V((unsigned) sel_stride), + mtl::Arg::V((int) gen), mtl::Arg::V((long) p0), mtl::Arg::V((int) t0), mtl::Arg::V((int) R), + mtl::Arg::V((int) kcap), mtl::Arg::B(bn), mtl::Arg::V((unsigned) cap), mtl::Arg::B(bo, off(ho)), + mtl::Arg::V((unsigned long) row_bytes)}; + d.grid[0] = (uint32_t) ((sel_stride + 15) / 16); + d.grid[1] = (uint32_t) T; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + g_cx->free(bn); + return true; +} + + +// ds4_kernels.cu:1772 - the expert outputs back into the token rows +bool moe_combine(void* acc, const void* y, const int* pos, const float* w_, int T, int k, int r0, int r1, int dim) { + if (T <= 0 || dim <= 0) return true; + if (!ensure()) return false; + void* ha = cpu_addr(acc); + void* hy = cpu_addr((void*) y); + void* hp = cpu_addr((void*) pos); + void* hw = cpu_addr((void*) w_); + if (!ha || !hy || !hp || !hw) { std::fprintf(stderr, "dsv41 metal: moe_combine has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer ba = wrap(ha, (size_t) T * dim * 4 ); + mtl::Buffer by = wrap(hy, (size_t) (r1 > r0 ? r1 - r0 : 1) * dim * 4 ); + mtl::Buffer bp = wrap(hp, (size_t) T * 8 * 4 ); + mtl::Buffer bw = wrap(hw, (size_t) T * 8 * 4 ); + if (!ba.handle || !by.handle || !bp.handle || !bw.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_moe_combine"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(ba, off(ha)), mtl::Arg::B(by, off(hy)), mtl::Arg::B(bp, off(hp)), mtl::Arg::B(bw, off(hw)), + mtl::Arg::V((int) k), mtl::Arg::V((int) r0), mtl::Arg::V((int) r1), mtl::Arg::V((unsigned) dim)}; + d.grid[0] = (uint32_t) ((dim + 255) / 256); + d.grid[1] = (uint32_t) T; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + + +// ds4_kernels.cu:362 - hc_post: out[t][k][e] = post[t][k]*y[t][e] + sum_j comb[t][j][k]*res[t][j][e] +bool hc_post_v(const float* y, const float* res, const float* post, const float* comb, int dim, int hc, int T, float* out) { + if (T <= 0 || dim <= 0) return true; + if (!ensure()) return false; + void* hy = cpu_addr((void*) y); + void* hr = cpu_addr((void*) res); + void* hp = cpu_addr((void*) post); + void* hco = cpu_addr((void*) comb); + void* ho = cpu_addr((void*) out); + if (!hy || !hr || !hp || !hco || !ho) { std::fprintf(stderr, "dsv41 metal: hc_post has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer by = wrap(hy, (size_t) T * dim * 4 ); + mtl::Buffer br = wrap(hr, (size_t) T * hc * dim * 4 ); + mtl::Buffer bp = wrap(hp, (size_t) T * 8 * 4 ); + mtl::Buffer bc = wrap(hco, (size_t) T * 64 * 4 ); + mtl::Buffer bo = wrap(ho, (size_t) T * hc * dim * 4 ); + if (!by.handle || !br.handle || !bp.handle || !bc.handle || !bo.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_hc_post_v"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(by, off(hy)), mtl::Arg::B(br, off(hr)), mtl::Arg::B(bp, off(hp)), + mtl::Arg::B(bc, off(hco)), mtl::Arg::B(bo, off(ho)), mtl::Arg::V((unsigned) dim), + mtl::Arg::V((unsigned) hc)}; + d.grid[0] = (uint32_t) ((dim * hc + 255) / 256); + d.grid[1] = (uint32_t) T; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + + +// ds4_kernels.cu:441 - the F32 GEMV (residual / dense projections) +bool gemv_f32(const float* W, const float* x, int n, int k, int T, float* y, int ldy) { + if (T <= 0 || n <= 0 || k <= 0) return true; + if (!ensure()) return false; + void* hw = cpu_addr((void*) W); + void* hx = cpu_addr((void*) x); + void* hy = cpu_addr((void*) y); + if (!hw || !hx || !hy) { std::fprintf(stderr, "dsv41 metal: gemv_f32 has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bw = wrap(hw, (size_t) n * k * 4 ); + mtl::Buffer bx = wrap(hx, (size_t) T * k * 4 ); + mtl::Buffer by = wrap(hy, (size_t) T * (ldy > 0 ? ldy : n) * 4 ); + if (!bw.handle || !bx.handle || !by.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_gemv_f32"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bw, off(hw)), mtl::Arg::B(bx, off(hx)), mtl::Arg::V((unsigned) n), + mtl::Arg::V((unsigned) k), mtl::Arg::B(by, off(hy)), mtl::Arg::V((unsigned) ldy)}; + d.grid[0] = (uint32_t) ((n + 7) / 8); + d.grid[1] = (uint32_t) T; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + + +// ds4_kernels.cu:1353 - the router's selection (idx, w, plan_ptr). The prefetch requests stay unpublished for now. +bool route_publish(const float* logits, const float* bias, int n, int k, int T, float route_scale, const void* vram, + int* idx, float* w, void* plan_ptr, const float* bias_vl, unsigned vl) { + if (T <= 0 || n <= 0) return true; + if (!ensure()) return false; + void* hl = cpu_addr((void*) logits); + void* hb = cpu_addr((void*) bias); + void* hv = cpu_addr((void*) vram); + void* hi = cpu_addr((void*) idx); + void* hw = cpu_addr((void*) w); + void* hp = cpu_addr((void*) plan_ptr); + if (!hl || !hb || !hv || !hi || !hw || !hp) { std::fprintf(stderr, "dsv41 metal: route_publish has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { // the pages the buffer itself spans (see q8_quant_rows' note) + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bl = wrap(hl, (size_t) T * 1024 * 4 ); + mtl::Buffer bb = wrap(hb, (size_t) n * 4 ); + mtl::Buffer bv = wrap(hv, (size_t) n * 8 ); + mtl::Buffer bi = wrap(hi, (size_t) T * 8 * 4 ); + mtl::Buffer bw = wrap(hw, (size_t) T * 8 * 4 ); + mtl::Buffer bp = wrap(hp, (size_t) T * 8 * 8 ); + if (!bl.handle || !bb.handle || !bv.handle || !bi.handle || !bw.handle || !bp.handle) return false; + mtl::Buffer bvl = bb; + size_t vl_off = 0; + if (bias_vl) { + void* hvl = cpu_addr((void*) bias_vl); + if (!hvl) return false; + bvl = wrap(hvl, (size_t) n * 4 ); + vl_off = off(hvl); + } + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_route_publish"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bl, off(hl)), mtl::Arg::B(bb, off(hb)), mtl::Arg::V((unsigned) n), mtl::Arg::V((unsigned) k), + mtl::Arg::V(route_scale), mtl::Arg::B(bv, off(hv)), mtl::Arg::B(bi, off(hi)), mtl::Arg::B(bw, off(hw)), + mtl::Arg::B(bp, off(hp)), mtl::Arg::B(bvl, vl_off), mtl::Arg::V((unsigned) vl), + mtl::Arg::V((unsigned) (bias_vl ? 1 : 0))}; + d.grid[0] = (uint32_t) T; + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 512; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + + +// ds4_kernels.cu swiglu: h = silu(g) * u (clamped when lim > 0) +bool swiglu(const float* g, const float* u, int n, float lim, float* h) { + if (n <= 0) return true; + if (!ensure()) return false; + void* hg = cpu_addr((void*) g); + void* hu = cpu_addr((void*) u); + void* hh = cpu_addr((void*) h); + if (!hg || !hu || !hh) { std::fprintf(stderr, "dsv41 metal: swiglu has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + (size_t) n * 4; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bg = wrap(hg), bu = wrap(hu), bh = wrap(hh); + if (!bg.handle || !bu.handle || !bh.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_swiglu"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bg, off(hg)), mtl::Arg::B(bu, off(hu)), mtl::Arg::V((unsigned) n), mtl::Arg::V(lim), + mtl::Arg::B(bh, off(hh))}; + d.grid[0] = (uint32_t) ((n + 255) / 256); + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + + +// vendor/moe_win.cu:49 moe_win_group - the expert -> group index for a window. One thread in CUDA; on the Mac every device +// pointer is a CPU address, so the same 25-line loop runs on the host, which keeps the first-appearance order the kernel has and +// avoids mirroring the 17-field MoeWin struct through a constant buffer. +bool moe_win_group_cpu(const void* plan_ptr, void* grp_ptr, void* grp_start, void* grp_ent, void* n_groups, int T, int k) { + if (T <= 0 || k <= 0) return true; + if (T * 8 > 4096) { std::fprintf(stderr, "dsv41 metal: moe_win_group window of %d too wide\n", T); return false; } + const unsigned long long* plan = (const unsigned long long*) cpu_addr((void*) plan_ptr); + unsigned long long* gptr = (unsigned long long*) cpu_addr(grp_ptr); + int* gstart = (int*) cpu_addr(grp_start); + int* gent = (int*) cpu_addr(grp_ent); + int* ng_out = (int*) cpu_addr(n_groups); + if (!plan || !gptr || !gstart || !gent || !ng_out) return false; + static unsigned long long gp[4096]; + static int cnt[4096], fill[4097], gof[4096]; + int ng = 0; + for (int t = 0; t < T; ++t) { + for (int i = 0; i < 8; ++i) { + const int e = t * 8 + i; + gof[e] = -1; + const unsigned long long p = i < k ? plan[e] : 0ull; + if (p == 0ull) continue; + int g = 0; + while (g < ng && gp[g] != p) ++g; + if (g == ng) { + gp[ng] = p; + cnt[ng++] = 0; + } + ++cnt[g]; + gof[e] = g; + } + } + fill[0] = 0; + for (int g = 0; g < ng; ++g) { + gptr[g] = gp[g]; + gstart[g] = fill[g]; + fill[g + 1] = fill[g] + cnt[g]; + } + gstart[ng] = fill[ng]; + for (int e = 0; e < T * 8; ++e) { + if (gof[e] >= 0) gent[fill[gof[e]]++] = e; + } + *ng_out = ng; + return true; +} + + + + +// ---- the prompt chunk's dense GEMM: float rows -> half, a quantized matrix -> half rows, MPP matmul2d (kernels/dense_mm.metal) +namespace { +mtl::Buffer g_gm_xs, g_gm_yt; // grow-only scratch: packed activations, the contiguous result +int g_lib_mm = -2; +mtl::Buffer scratch(mtl::Buffer& b, size_t bytes) { + if (!b.handle || b.bytes < bytes) { + if (b.handle) g_cx->free(b); + b = g_cx->alloc(bytes + (bytes >> 2)); + } + return b; +} +} // namespace + +bool rows_f16(const float* x, int ldx, const int* idx, int nrows, int ncols, uint16_t* out) { + if (nrows <= 0 || ncols <= 0) return true; + if (!ensure()) return false; + void* hx = strata::gpu::host_view((void*) x); + void* ho = strata::gpu::host_view(out); + void* hi = idx ? strata::gpu::host_view((void*) idx) : nullptr; + if (!hx || !ho) return false; + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bx = wrap(hx, ((size_t) (nrows - 1) * (size_t) ldx + (size_t) ncols) * 4); + mtl::Buffer bo = wrap(ho, (size_t) nrows * (size_t) ncols * 2); + mtl::Buffer bi = hi ? wrap(hi, (size_t) nrows * 4) : bx; + if (!bx.handle || !bo.handle || !bi.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_rows_f16"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bx, off(hx)), mtl::Arg::V((unsigned long) ldx), mtl::Arg::B(bi, hi ? off(hi) : off(hx)), + mtl::Arg::V((unsigned) (hi ? 1 : 0)), mtl::Arg::V((unsigned) ncols), mtl::Arg::B(bo, off(ho))}; + d.grid[0] = (uint32_t) ((ncols + 255) / 256); + d.grid[1] = (uint32_t) nrows; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + return true; +} + +// kind: 0 BF16 values, 8 Q8_0, 13 Q5_K, 14 Q6_K. n rows of k values, ld the row stride (bytes; values for BF16), out half rows. +bool dequant_rows(int kind, const void* raw, size_t ld, int k, int n, uint16_t* out) { + if (n <= 0 || k <= 0) return true; + if (!ensure()) return false; + const char* fn = kind == 13 ? "ds4_dq_q5k" : kind == 14 ? "ds4_dq_q6k" : kind == 8 ? "ds4_dq_q8" : kind == 0 ? "ds4_dq_bf16" : nullptr; + if (!fn) return false; + if (((kind == 13 || kind == 14) && k % 256 != 0) || (kind == 8 && k % 32 != 0)) return false; + void* hw = strata::gpu::host_view((void*) raw); + void* ho = strata::gpu::host_view(out); + if (!hw || !ho) return false; + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + const size_t wbytes = (size_t) (n - 1) * ld * (kind == 0 ? 2 : 1) + (kind == 0 ? (size_t) k * 2 : (kind == 13 ? (size_t) (k / 256) * 176 : kind == 14 ? (size_t) (k / 256) * 210 : (size_t) (k / 32) * 34)); + if (!std::getenv("DS4_NO_TOUCH")) { // (the matrix may be file-backed pages not read yet: the GPU does not page them in, it stalls) + volatile unsigned char sink = 0; + for (size_t o = 0; o < wbytes; o += 4096) sink = sink + ((const volatile unsigned char*) hw)[o]; + (void) sink; + } + mtl::Buffer bw = wrap(hw, wbytes); + mtl::Buffer bo = wrap(ho, (size_t) n * (size_t) k * 2); + if (!bw.handle || !bo.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, fn); } catch (...) { return false; } + size_t items = kind == 13 ? (size_t) n * (k / 256) * 8 : kind == 14 ? (size_t) n * (k / 256) * 64 : kind == 8 ? (size_t) n * (k / 32) : (size_t) n * k; + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bw, off(hw)), mtl::Arg::V((unsigned long) ld), mtl::Arg::V((unsigned) k), mtl::Arg::V((unsigned) n), + mtl::Arg::B(bo, off(ho))}; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + d.grid[0] = (uint32_t) ((items + 255) / 256); + d.grid[1] = d.grid[2] = 1; + g_cx->run({d}); + return true; +} + +// Y[r][c] = alpha * (X16[r] . W16[c]) + beta * Y[r][c]: X16 rows x K (stride ldx halves), W16 N x K contiguous, Y rows x N float +// (stride ldy). The cuBLAS call's operands in row-major terms; the MPP tile kernel wants contiguous operands and result. +bool gemm_f16_nt(const uint16_t* W16, const uint16_t* X16, int ldx, float* Y, int ldy, int rows, int N, int K, float alpha, + float beta) { + if (rows <= 0 || N <= 0) return true; + if (!ensure()) return false; + if (g_lib_mm == -2) { + g_lib_mm = g_cx->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal"}); + if (g_lib_mm < 0) std::fprintf(stderr, "dsv41 metal: compiling dense_mm.metal failed\n"); + } + if (g_lib_mm < 0) return false; + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + void* hw = strata::gpu::host_view((void*) W16); + void* hx = strata::gpu::host_view((void*) X16); + void* hy = strata::gpu::host_view(Y); + if (!hw || !hx || !hy) return false; + mtl::Buffer bw = wrap(hw, (size_t) N * K * 2); + mtl::Buffer bx = wrap(hx, ((size_t) (rows - 1) * (size_t) ldx + (size_t) K) * 2); + mtl::Buffer by = wrap(hy, ((size_t) (rows - 1) * (size_t) ldy + (size_t) N) * 4); + if (!bw.handle || !bx.handle || !by.handle) return false; + // activations contiguous + mtl::Buffer ax = bx; + size_t ax_off = off(hx); + if (ldx != K) { + ax = scratch(g_gm_xs, (size_t) rows * K * 2); + ax_off = 0; + mtl::Pipeline pp; + try { pp = g_cx->pipeline(g_lib, "ds4_pack_h16"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pp; + d.args = {mtl::Arg::B(bx, off(hx)), mtl::Arg::V((unsigned long) ldx), mtl::Arg::V((unsigned) K), mtl::Arg::B(ax)}; + d.grid[0] = (uint32_t) ((K + 255) / 256); + d.grid[1] = (uint32_t) rows; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + } + const bool direct = ldy == N && alpha == 1.0f && beta == 0.0f; + mtl::Buffer ty = by; + size_t ty_off = off(hy); + if (!direct) { + ty = scratch(g_gm_yt, (size_t) rows * N * 4); + ty_off = 0; + } + const int tm = 64, tn = N <= 64 ? 32 : 64; + const std::string fn = "gemm_f16_" + std::to_string(tm) + "_" + std::to_string(tn) + "_4"; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib_mm, fn); } catch (...) { return false; } + struct GemmArgs { int M, N, K, pad; } ga{rows, N, K, 0}; + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(ax, ax_off), mtl::Arg::B(bw, off(hw)), mtl::Arg::B(ty, ty_off), mtl::Arg::V(ga)}; + d.grid[0] = (uint32_t) ((N + tn - 1) / tn); + d.grid[1] = (uint32_t) ((rows + tm - 1) / tm); + d.grid[2] = 1; + d.tg[0] = 128; d.tg[1] = d.tg[2] = 1; + g_cx->run({d}); + if (!direct) { + mtl::Pipeline pa; + try { pa = g_cx->pipeline(g_lib, "ds4_axpby_rows"); } catch (...) { return false; } + mtl::Dispatch e; + e.p = &pa; + e.args = {mtl::Arg::B(ty), mtl::Arg::B(by, off(hy)), mtl::Arg::V((unsigned long) ldy), mtl::Arg::V((unsigned) N), + mtl::Arg::V(alpha), mtl::Arg::V(beta)}; + e.grid[0] = (uint32_t) ((N + 255) / 256); + e.grid[1] = (uint32_t) rows; + e.grid[2] = 1; + e.tg[0] = 256; e.tg[1] = e.tg[2] = 1; + g_cx->run({e}); + } + return true; +} + + +// ---- the prompt chunk's expert GEMMs (ds4_mm.metal): a set = n_exp consecutive expert slots, their rows sorted by expert ---- +namespace { +int g_lib_moe = -2; +mtl::Buffer g_mm_xs, g_mm_hs, g_mm_ys; // grow-only scratch: the sorted rows' half activations; the gate/up's swiglu rows (the down's input) +bool g_mm_hs_ok = false; +// bytes of a 256-value block per type: the tile kernels are instantiated for exactly these (DS4_MM_INST) +size_t iq_bs(int type) { return type == 16 ? 66 : type == 17 ? 74 : type == 22 ? 82 : type == 18 ? 98 : type == 29 ? 56 : type == 19 ? 50 : 0; } +struct DsMmArgsH { + unsigned long gate_off, up_off, down_off, stride; + unsigned gu_row, d_row, n_embd, n_ff; + float limit; + unsigned pad0; +}; +bool mm_prepare() { + if (!ensure()) return false; + if (g_lib_moe == -2) { + g_lib_moe = g_cx->compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal", "../src/engine/msl/ds4_mm.metal"}); + if (g_lib_moe < 0) std::fprintf(stderr, "dsv41 metal: compiling ds4_mm.metal failed\n"); + } + return g_lib_moe >= 0; +} +} // namespace + +static bool mm_set(bool gate_up, const uint8_t* wbase, size_t stride, const size_t* tr, const int* hb, int n_exp, int total_rows, + int n_embd, int n_ff, float limit, const float* xf, int ldx, const int* idx, const uint16_t* hin, uint16_t* hout, + int ldh, float* yout, int ldy, int tg, int td) { + const int type = gate_up ? tg : td; + if (n_exp <= 0 || total_rows <= 0) return true; + if (iq_bs(tg) == 0 || iq_bs(td) == 0 || n_embd % 256 != 0 || n_ff % 256 != 0) return false; + if (ldh != n_ff || (!gate_up && ldy != n_embd)) return false; + if (gate_up ? (n_ff % 128 != 0) : (n_embd % 128 != 0)) return false; + if (!mm_prepare()) return false; + constexpr int TM = 32; // (the tile kernels' TM: ds4_mm.metal DS4_MM_INST) + constexpr int SG = 4, TN = 32; // the tile kernels' simdgroups and columns each: + constexpr int COLS = SG * TN; // ds4_mm.metal DS4_MM_INST must match these + // the tile list: (slot, first row, valid rows) a TM-row piece of one expert + std::vector tiles; + for (int e = 0; e < n_exp; ++e) { + const int lo = hb[e] - hb[0], hi = hb[e + 1] - hb[0]; + for (int r = lo; r < hi; r += TM) { tiles.push_back(e); tiles.push_back(r); tiles.push_back(std::min(TM, hi - r)); tiles.push_back(0); } + } + const int ntiles = (int) (tiles.size() / 4); + if (ntiles == 0) return true; + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + const size_t len = ((end - base + page - 1) / page) * page; + if (std::getenv("DS4_MM_DBG")) { size_t o2 = 0; std::fprintf(stderr, "mm wrap %p len %zu own=%d rem=%zu\n", (void*) base, len, ds4_gpu_buffer_for_cpu((void*) base, len, &o2) != nullptr, ds4_gpu_alloc_remaining((void*) base)); } + return g_cx->wrap((void*) base, len); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + void* hw = strata::gpu::host_view((void*) wbase); + if (!hw) return false; + mtl::Buffer bw = wrap(hw, (size_t) n_exp * stride); + mtl::Buffer bt = g_cx->alloc(tiles.size() * 4); + std::memcpy(bt.ptr, tiles.data(), tiles.size() * 4); + DsMmArgsH a{tr[0], tr[1], tr[2], (unsigned long) stride, (unsigned) ((size_t) n_embd / 256 * iq_bs(tg)), (unsigned) ((size_t) n_ff / 256 * iq_bs(td)), + (unsigned) n_embd, (unsigned) n_ff, limit, std::getenv("DS4_MM_ABL") ? (unsigned) std::atoi(std::getenv("DS4_MM_ABL")) : 0u}; + if (std::getenv("DS4_MM_DBG")) std::fprintf(stderr, "mm %s embd %d ff %d tg %d td %d lim %g stride %zu tr %zu %zu %zu gu_row %u d_row %u\n", gate_up ? "gu" : "dn", n_embd, n_ff, tg, td, limit, stride, tr[0], tr[1], tr[2], a.gu_row, a.d_row); + mtl::Pipeline pl; + // DS4_MM_CTAB: QUANT_CTAB (function constant 23) - the i-quant grids and the sign table read from constant + // memory instead of a threadgroup copy. Same values; the decode's two lookups per 8 weights go through the + // constant cache instead of threadgroup memory. + static const bool mm_ctab = std::getenv("DS4_MM_CTAB") != nullptr; + try { + const std::string fn = std::string(gate_up ? "ds4_mm_gu_" : "ds4_mm_dn_") + std::to_string(type); + if (mm_ctab) { + mtl::Constants cs; + cs.b(23, true); + pl = g_cx->pipeline_uncached(g_lib_moe, fn, cs); + } else { + pl = g_cx->pipeline(g_lib_moe, fn); + } + } catch (...) { g_cx->free(bt); return false; } + if (std::getenv("DS4_MM_DBG")) std::fprintf(stderr, "pipe %s max_threads %d tgmem %d width %d\n", pl.name.c_str(), pl.max_threads, pl.static_tg_mem, pl.simd_width); + mtl::Dispatch d; + d.p = &pl; + d.tg[0] = SG * 32; d.tg[1] = d.tg[2] = 1; + d.grid[2] = 1; + d.grid[1] = (uint32_t) ntiles; + if (gate_up) { + // the sorted rows' activations, half + const size_t xs_bytes = ((size_t) total_rows + TM) * (size_t) n_embd * 2; + if (!g_mm_xs.handle || g_mm_xs.bytes < xs_bytes) { + if (g_mm_xs.handle) g_cx->free(g_mm_xs); + g_mm_xs = g_cx->alloc(xs_bytes + (xs_bytes >> 2)); + } + void* hx = strata::gpu::host_view((void*) xf); + void* hi = strata::gpu::host_view((void*) idx); + void* ho = strata::gpu::host_view(hout); + if (!hx || !hi || !ho) { g_cx->free(bt); return false; } + // (the float rows: the whole token slab the engine's c.x holds; its extent is the allocation's) + const size_t xrem = ds4_gpu_alloc_remaining(hx); + mtl::Buffer bxf = wrap(hx, xrem > 0 ? xrem : (size_t) ldx * 4); + mtl::Buffer bix = wrap(hi, (size_t) total_rows * 4); + mtl::Pipeline pg; + try { pg = g_cx->pipeline(g_lib_moe, "ds4_mm_gather"); } catch (...) { g_cx->free(bt); return false; } + mtl::Dispatch dg; + dg.p = &pg; + dg.args = {mtl::Arg::B(bxf, off(hx)), mtl::Arg::V((unsigned long) ldx), mtl::Arg::B(bix, off(hi)), mtl::Arg::V((unsigned) n_embd), mtl::Arg::B(g_mm_xs)}; + dg.grid[0] = (uint32_t) ((n_embd + 255) / 256); + dg.grid[1] = (uint32_t) total_rows; + dg.grid[2] = 1; + dg.tg[0] = 256; dg.tg[1] = dg.tg[2] = 1; + g_cx->run({dg}); + const size_t hs_bytes = ((size_t) total_rows + TM) * (size_t) n_ff * 2; + if (!g_mm_hs.handle || g_mm_hs.bytes < hs_bytes) { + if (g_mm_hs.handle) g_cx->free(g_mm_hs); + g_mm_hs = g_cx->alloc(std::max(hs_bytes + (hs_bytes >> 2), std::getenv("DS4_MM_BIG") ? (size_t) 1 << 26 : 0)); + } + g_mm_hs_ok = true; + { // (the down's result rows are made here too) + const size_t ys_bytes = ((size_t) total_rows + TM) * (size_t) n_embd * 4; + if (!g_mm_ys.handle || g_mm_ys.bytes < ys_bytes) { + if (g_mm_ys.handle) g_cx->free(g_mm_ys); + g_mm_ys = g_cx->alloc(std::max(ys_bytes + (ys_bytes >> 2), std::getenv("DS4_MM_BIG") ? (size_t) 1 << 26 : 0)); + } + } + (void) ho; + d.args = {mtl::Arg::B(bw, off(hw)), mtl::Arg::B(g_mm_xs), mtl::Arg::B(g_mm_hs), mtl::Arg::B(bt), mtl::Arg::V(a), mtl::Arg::V((unsigned) ntiles)}; + d.grid[0] = (uint32_t) (n_ff / COLS); + } else { + void* hh = strata::gpu::host_view((void*) hin); + void* hy = strata::gpu::host_view(yout); + void* hy_real = nullptr; + if (!hh || !hy) { g_cx->free(bt); return false; } + const size_t hrem = ds4_gpu_alloc_remaining(hh); + if (!g_mm_hs_ok) { g_cx->free(bt); return false; } + g_mm_hs_ok = false; + (void) hrem; + mtl::Buffer& bh = g_mm_hs; + const size_t ys_bytes = ((size_t) total_rows + TM) * (size_t) n_embd * 4; + if (!g_mm_ys.handle || g_mm_ys.bytes < ys_bytes) { + if (g_mm_ys.handle) g_cx->free(g_mm_ys); + g_mm_ys = g_cx->alloc(std::max(ys_bytes + (ys_bytes >> 2), std::getenv("DS4_MM_BIG") ? (size_t) 1 << 26 : 0)); + } + mtl::Buffer by = g_mm_ys; + hy_real = hy; + hy = by.ptr; + if (std::getenv("DS4_MM_YALT")) { hy = hh; by = wrap(hh, (size_t) total_rows * (size_t) n_ff * 2); } + if (std::getenv("DS4_MM_EXP1")) { by = g_cx->alloc((size_t) total_rows * (size_t) n_embd * 4 + 4096); hy = by.ptr; } + d.args = {mtl::Arg::B(bw, off(hw)), mtl::Arg::B(g_mm_ys), mtl::Arg::B(g_mm_hs), mtl::Arg::B(bt), mtl::Arg::V(a), mtl::Arg::V((unsigned) ntiles)}; + d.grid[0] = (uint32_t) (n_embd / COLS); + } + if (!gate_up && std::getenv("DS4_MM_ASGU")) { // test: the gate/up kernel on the down's buffers + static mtl::Pipeline pgu; + pgu = g_cx->pipeline(g_lib_moe, "ds4_mm_gu_" + std::to_string(tg)); + d.p = &pgu; + d.args[1] = mtl::Arg::B(g_mm_xs); + d.args[2] = mtl::Arg::B(g_mm_ys); + if (std::atoi(std::getenv("DS4_MM_ASGU")) == 1) d.grid[0] = (uint32_t) (n_ff / COLS); + } + if (!gate_up && std::getenv("DS4_MM_DROP")) { + const char* dk = std::getenv("DS4_MM_DROP"); + static mtl::Buffer dummy; + if (!dummy.handle) dummy = g_cx->alloc(1 << 26); + static mtl::Buffer dummy2; + if (!dummy2.handle) dummy2 = g_cx->alloc(1 << 26); + for (int k = 0; k < 4; ++k) if (std::strchr(dk, '0' + k)) d.args[k] = mtl::Arg::B(std::getenv("DS4_MM_D2") && k == std::atoi(std::getenv("DS4_MM_D2")) ? dummy2 : dummy); + } + if (!gate_up && std::getenv("DS4_MM_ONE")) { d.grid[1] = 1; if (std::getenv("DS4_MM_DBG")) std::fprintf(stderr, "dn tile0 %d %d %d ntiles %d x %u rows %d\n", tiles[0], tiles[1], tiles[2], ntiles, d.grid[0], total_rows); } + if (std::getenv("DS4_MM_DBG")) for (size_t q = 0; q < d.args.size(); ++q) if (d.args[q].buf) std::fprintf(stderr, "arg %s %zu bytes %zu bias %zu off %zu ptr %p\n", gate_up ? "gu" : "dn", q, d.args[q].buf->bytes, (size_t) d.args[q].buf->bias, (size_t) d.args[q].off, d.args[q].buf->ptr); + if (std::getenv(gate_up ? "DS4_MM_NOGU" : "DS4_MM_NODN")) { g_cx->free(bt); return true; } + g_cx->run({d}); + if (!std::getenv("DS4_MM_LEAK")) g_cx->free(bt); + if (!gate_up) { // the result into the engine's rows + void* hyr = strata::gpu::host_view(yout); + const size_t n4 = (size_t) total_rows * (size_t) n_embd; + const size_t base = (size_t) ((uintptr_t) hyr) & ~(size_t) (page - 1); + mtl::Buffer bd = g_cx->wrap((void*) base, wrap_len(base, (size_t) (uintptr_t) hyr + n4 * 4 - base)); + mtl::Pipeline pc; + try { pc = g_cx->pipeline(g_lib, "ds4_copy4"); } catch (...) { return false; } + mtl::Dispatch c; + c.p = &pc; + c.args = {mtl::Arg::B(g_mm_ys), mtl::Arg::B(bd, (size_t) ((uintptr_t) hyr - base)), mtl::Arg::V((unsigned long) n4)}; + c.tg[0] = 256; c.tg[1] = c.tg[2] = 1; + c.grid[0] = (uint32_t) ((n4 + 255) / 256); + c.grid[1] = c.grid[2] = 1; + g_cx->run({c}); + } + return true; +} + +bool iq_set_gu(const uint8_t* wbase, size_t stride, const size_t* tr, const int* hbounds, int n_exp, int total_rows, int n_embd, int n_ff, + float limit, const float* xf, int ldx, const int* idx, uint16_t* h, int ldh, int tg, int td) { + return mm_set(true, wbase, stride, tr, hbounds, n_exp, total_rows, n_embd, n_ff, limit, xf, ldx, idx, nullptr, h, ldh, nullptr, 0, tg, td); +} +bool iq_set_dn(const uint8_t* wbase, size_t stride, const size_t* tr, const int* hbounds, int n_exp, int total_rows, int n_embd, int n_ff, + float limit, const uint16_t* h16, int ldh, float* y, int ldy, int tg, int td) { + return mm_set(false, wbase, stride, tr, hbounds, n_exp, total_rows, n_embd, n_ff, limit, nullptr, 0, nullptr, h16, nullptr, ldh, y, ldy, tg, td); +} + +// ---- copies, fills and event marks on the context's queue (see ds4_copy4 in ds4_init.metal) ------------------------------- +// strata::gpu routes cudaMemcpyAsync / cudaMemsetAsync here when both ends are its allocations; ds4_ctx_mark ends the batched +// command buffer without waiting and returns a token the matching ds4_ctx_wait blocks on (cudaEventRecord / EventSynchronize). +// DS4_COPY_LOG: how much host time the copies take (they wrap host memory into Metal buffers per call) +namespace { struct CopyLog { unsigned long long n = 0, bytes = 0, us = 0; + ~CopyLog() { if (n) std::fprintf(stderr, "copy log: %llu copies, %.1f MB, %.1f ms host time (%.1f us a copy)\n", + n, bytes / 1e6, us / 1e3, n ? (double) us / (double) n : 0.0); } }; + CopyLog g_copy_us; } +static inline unsigned long long copy_now_us() { + return (unsigned long long) std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()).count(); +} +extern "C" bool ds4_ctx_copy(void* dst, const void* src, size_t n) { + if (n == 0) return true; + const bool copy_log = std::getenv("DS4_COPY_LOG") != nullptr; + const unsigned long long copy_t0 = copy_log ? copy_now_us() : 0; + struct CopyDone { bool on; unsigned long long t0; size_t n; ~CopyDone() { if (on) { g_copy_us.n++; g_copy_us.bytes += n; g_copy_us.us += copy_now_us() - t0; } } } copy_done{copy_log, copy_t0, n}; + if (!ensure()) return false; + void* hd = strata::gpu::host_view(dst); + void* hs = strata::gpu::host_view((void*) src); + if (!hd || !hs) return false; + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + const size_t base = (size_t) ((uintptr_t) host) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) host) + bytes; + return g_cx->wrap((void*) base, wrap_len(base, end - base)); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer bs = wrap(hs, n); + mtl::Buffer bd = wrap(hd, n); + if (!bs.handle || !bd.handle) return false; + const bool w4 = (((uintptr_t) hd | (uintptr_t) hs | n) & 3u) == 0; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, w4 ? "ds4_copy4" : "ds4_copy1"); } catch (...) { return false; } + const size_t items = w4 ? n / 4 : n; + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bs, off(hs)), mtl::Arg::B(bd, off(hd)), mtl::Arg::V((unsigned long) items)}; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + d.grid[0] = (uint32_t) ((items + 255) / 256); + d.grid[1] = d.grid[2] = 1; + g_cx->run({d}); + return true; +} +extern "C" bool ds4_ctx_fill(void* dst, int value, size_t n) { + if (n == 0) return true; + if (!ensure()) return false; + void* hd = strata::gpu::host_view(dst); + if (!hd) return false; + const size_t page = 16384; + const size_t base = (size_t) ((uintptr_t) hd) & ~(size_t) (page - 1); + const size_t end = (size_t) ((uintptr_t) hd) + n; + mtl::Buffer bd = g_cx->wrap((void*) base, wrap_len(base, end - base)); + if (!bd.handle) return false; + mtl::Pipeline pl; + try { pl = g_cx->pipeline(g_lib, "ds4_fill1"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(bd, (size_t) ((uintptr_t) hd - base)), mtl::Arg::V((unsigned long) n), mtl::Arg::V((unsigned) (value & 0xff))}; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + d.grid[0] = (uint32_t) ((n + 255) / 256); + d.grid[1] = d.grid[2] = 1; + g_cx->run({d}); + return true; +} +extern "C" void* ds4_ctx_mark() { + if (!g_cx) return nullptr; + return g_cx->mark(); +} +extern "C" void ds4_ctx_drop(void* token) { + if (g_cx && token) g_cx->drop_mark(token); +} +extern "C" void ds4_ctx_wait(void* token) { + static const bool no_wait_split = std::getenv("DS4_NO_WAIT_SPLIT") != nullptr; + if (g_cx && !no_wait_split) g_cx->split_runs(); // (what is queued after the mark starts on the GPU while the host waits) + if (g_cx && token) g_cx->wait_mark(token); +} +// commit the batched dispatches without waiting: the GPU starts on them while the host goes on (the CPU lane) +extern "C" void ds4_metal_kick() { + if (g_cx) g_cx->split_runs(); +} + +extern "C" void ds4_metal_flush() { + // The engine's copies (the embedding rows, the streamed experts, the readbacks) are encoded on strata::gpu's default + // stream, its dispatched kernels on mtl_rt's deferred command buffer. Two queues, so committing one alone leaves the + // other's work unordered: a kernel could read a feed whose blit had not run yet (emb_rows is a zeroed buffer, so + // embed wrote exactly 0.0 into h). The stream's queue goes first - it only carries the feeds at a token boundary, + // the readbacks are encoded after this returns. + ds4_gpu_sync_all_streams(); + if (g_cx) g_cx->flush_runs(); +} + + +// the one Context every dsv41 kernel dispatches through: two of them means two Metal queues and no ordering between their work +mtl::Context& shared_cx() { + ensure(); + return *g_cx; +} +mtl::Context& finish_cx() { + ensure(); + return *g_cx; +} + +int shared_lib() { + ensure(); + return g_lib; +} + +} // namespace strata::dsv41::metal_host \ No newline at end of file diff --git a/edge0-strata/metal/src/engine/ds4_host.cpp b/edge0-strata/metal/src/engine/ds4_host.cpp new file mode 100644 index 00000000..2d425174 --- /dev/null +++ b/edge0-strata/metal/src/engine/ds4_host.cpp @@ -0,0 +1,254 @@ +// metal/src/engine/ds4_host.cpp - host side of the ported dsv41 kernels: compile the MSL once, wrap the engine's own buffers, +// dispatch. First entry point: the Q8 activation quantiser (ds4_gguf.cu:201). +// +// The two runtimes meet on pointers: strata::gpu owns the engine's device memory (a device pointer is its MTLBuffer's +// gpuAddress) and exposes host_view(device_pointer); mtl::Context::wrap() adopts that address without copying and a +// Dispatch names it through Arg::B(). See RIGS/m1002.md for the derivation. +#include "mtl_rt.h" + +#include + + + +#include +#include +#include +#include + +#include +#include + +namespace strata::dsv41::metal_host { +mtl::Context& finish_cx(); // ds4_finish.cpp's Context: every kernel that orders with the chunk path's dispatches uses it +int shared_lib(); // from ds4_finish.cpp: the library the shared Context compiled + + +namespace { +mtl::Context* g_cx = nullptr; +int g_lib = -1; +bool g_tried = false; + +bool ensure(mtl::Context& cx) { + if (g_cx == &cx && g_lib >= 0) return true; + if (g_tried && g_lib < 0) return false; + g_tried = true; + // the kernel lives in its own file and is self-contained (metal_stdlib + ggml_grids.h) + g_lib = cx.compile({"../src/engine/msl/ds4_quant.metal"}, "", false); + if (g_lib < 0) { + std::fprintf(stderr, "dsv41 metal: compiling ds4_quant.metal failed\n"); + return false; + } + g_cx = &cx; + return true; +} +} // namespace + +// q8(x, ldx, n) -> q (int8, n bytes) and d (half, n/32); ldx is the stride in values between blocks of 32 +bool q8_quant(mtl::Context& cx, const float* x, int ldx, int n, void* q, void* d) { + if (!ensure(cx)) return false; + const size_t xb = (size_t) ldx * (size_t) (n / 32) * sizeof(float), qb = (size_t) n, db = (size_t) (n / 32) * 2; + mtl::Buffer bx = cx.wrap(strata::gpu::host_view(x), xb); + mtl::Buffer bq = cx.wrap(q, qb); + mtl::Buffer bd = cx.wrap(d, db); + mtl::Pipeline p; + try { + p = cx.pipeline(g_lib, "ds4_q8_quant"); + } catch (...) { + return false; + } + mtl::Dispatch ds; + ds.p = &p; + ds.args = {mtl::Arg::B(bx), mtl::Arg::V((unsigned) ldx), mtl::Arg::B(bq), mtl::Arg::B(bd)}; + ds.grid[0] = (uint32_t) (n / 32); + ds.grid[1] = ds.grid[2] = 1; + ds.tg[0] = 32; + ds.tg[1] = ds.tg[2] = 1; + mtl::require_fit(p, 32, "ds4_q8_quant"); + cx.run({ds}); + return true; +} + +} // namespace strata::dsv41::metal_host + +namespace strata::dsv41::metal_host { +namespace { +mtl::Context* g_init_cx = nullptr; +int g_lib_init = -1; + +mtl::Context& shared_cx() { + if (!g_init_cx) { + g_init_cx = new mtl::Context(); + g_lib_init = g_init_cx->compile({"../src/engine/msl/ds4_init.metal"}, "", false); + if (g_lib_init < 0) std::fprintf(stderr, "dsv41 metal: compiling ds4_init.metal failed\n"); + } + return *g_init_cx; +} +} // namespace + +// ds4_kernels.cu:1477 - pre[i] = (i % 8) == 0 ? 1 : 0 for i < T*8 +bool init_pre(float* pre, int T, void*) { + // cpu path (T*8 floats of a fixed pattern): the GPU route had to wrap a whole page to reach 32 bytes, and that aliasing + // buffer overlapped the pages holding h + void* host = strata::gpu::host_view(pre); + if (!host) { std::fprintf(stderr, "dsv41 metal: init_pre has no host view\n"); return false; } + float* p = (float*) host; + for (int t = 0; t < T; ++t) + for (int i = 0; i < 8; ++i) p[t * 8 + i] = (i == 0) ? 1.0f : 0.0f; + return true; +} + + +// ds4_kernels.cu:130 - the BF16 embedding gather, hc copies a token +bool embed(const uint16_t* rows, int T, float* h, int dim, int hc, void*) { + strata::dsv41::metal_host::shared_cx(); + if (g_lib_init < 0) { std::fprintf(stderr, "dsv41 metal: embed: no library\n"); return false; } + // host_view() returns a different address from the device pointer on this runtime, so wrapping it wrote into a mapping the + // engine never reads. On Apple silicon a device pointer is a valid CPU address for shared storage: wrap that instead. + // DS4_HOSTVIEW=1 restores the old behaviour. + const bool use_hv = std::getenv("DS4_HOSTVIEW") != nullptr; + void* hrows = use_hv ? strata::gpu::host_view((void*) rows) : (void*) rows; + void* hh = use_hv ? strata::gpu::host_view(h) : (void*) h; + if (!hrows || !hh) { std::fprintf(stderr, "dsv41 metal: embed has no host view\n"); return false; } + if (std::getenv("DS4_DBG_HV")) + std::fprintf(stderr, "DBG hostview: rows %p -> %p h %p -> %p\n", + (void*) rows, hrows, (void*) h, hh); + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + return strata::dsv41::metal_host::shared_cx().wrap((void*) (((uintptr_t) host) & ~(uintptr_t) (page - 1)), ((bytes + page - 1) / page) * page); + }; + mtl::Buffer br = wrap(hrows, (size_t) T * dim * 2); + mtl::Buffer bh = wrap(hh, (size_t) T * hc * dim * 4); + if (!br.handle || !bh.handle) { std::fprintf(stderr, "dsv41 metal: embed wrap failed\n"); return false; } + const size_t dro = (size_t) (((uintptr_t) hrows) - (((uintptr_t) hrows) & ~(uintptr_t) (page - 1))); + const size_t dho = (size_t) (((uintptr_t) hh) - (((uintptr_t) hh) & ~(uintptr_t) (page - 1))); + mtl::Pipeline pl; + try { pl = strata::dsv41::metal_host::shared_cx().pipeline(g_lib_init, "ds4_embed"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(br, dro), mtl::Arg::B(bh, dho), mtl::Arg::V((unsigned) dim), mtl::Arg::V((unsigned) hc)}; + d.grid[0] = (uint32_t) ((dim + 255) / 256); + d.grid[1] = (uint32_t) T; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + strata::dsv41::metal_host::shared_cx().run({d}); + if (std::getenv("DS4_DBG_EMBWRITE")) { + strata::dsv41::metal_host::shared_cx().flush_runs(); + const float* f0 = (const float*) hh; + const uint16_t* r16 = (const uint16_t*) hrows; + std::fprintf(stderr, "DBG embwrite: h=%p rows=%p h_off=%zu rows_off=%zu\n", hh, hrows, dho, dro); + std::fprintf(stderr, "DBG embwrite: rows[0..3]=%04x %04x %04x %04x h[0..3]=%g %g %g %g\n", + (unsigned) r16[0], (unsigned) r16[1], (unsigned) r16[2], (unsigned) r16[3], + f0[0], f0[1], f0[2], f0[3]); + } + return true; +} + + +// ds4_model.cu:2737 - C[t][m] = sum_k A[m][k] * B[t][k], F32 operands, alpha = 1, beta = 0 +bool sgemm_tn(const float* A, const float* B, float* C, int m, int nt, int k, int lda, int ldb, int ldc) { + finish_cx(); + const int lib = shared_lib(); + if (lib < 0) return false; + void* ha = strata::gpu::host_view((void*) A); + void* hb = strata::gpu::host_view((void*) B); + void* hc = strata::gpu::host_view((void*) C); + if (!ha || !hb || !hc) { std::fprintf(stderr, "dsv41 metal: sgemm_tn has no host view\n"); return false; } + const size_t page = 16384; + auto wrap = [&](void* host, size_t bytes) { + return finish_cx().wrap((void*) (((uintptr_t) host) & ~(uintptr_t) (page - 1)), ((bytes + page - 1) / page) * page); + }; + auto off = [&](void* host) { return (size_t) ((uintptr_t) host - (((uintptr_t) host) & ~(uintptr_t) (page - 1))); }; + mtl::Buffer ba = wrap(ha, (size_t) m * lda * 4); + mtl::Buffer bb = wrap(hb, (size_t) nt * ldb * 4); + mtl::Buffer bc = wrap(hc, (size_t) nt * ldc * 4); + if (!ba.handle || !bb.handle || !bc.handle) return false; + mtl::Pipeline pl; + try { pl = finish_cx().pipeline(lib, "ds4_sgemm_tn"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(ba, off(ha)), mtl::Arg::B(bb, off(hb)), mtl::Arg::B(bc, off(hc)), + mtl::Arg::V((unsigned) m), mtl::Arg::V((unsigned) nt), mtl::Arg::V((unsigned) k), + mtl::Arg::V((unsigned) lda), mtl::Arg::V((unsigned) ldb), mtl::Arg::V((unsigned) ldc)}; + const uint64_t total = (uint64_t) m * (uint64_t) nt; + d.grid[0] = (uint32_t) ((total + 127) / 128); + d.grid[1] = d.grid[2] = 1; + d.tg[0] = 128; d.tg[1] = d.tg[2] = 1; + finish_cx().run({d}); + if (std::getenv("DS4_DBG_SG")) { + finish_cx().flush_runs(); + const float* fa = (const float*) ha; const float* fb = (const float*) hb; const float* fc = (const float*) hc; + std::fprintf(stderr, "sgemm m %d nt %d k %d lda %d ldb %d ldc %d A %g %g B %g %g C %g %g %g\n", m, nt, k, lda, ldb, ldc, fa[0], fa[1], fb[0], fb[1], fc[0], fc[1], fc[2]); + } + return true; +} + + +// embed through raw device addresses: the struct carries the gpuAddress of rows and h, which the runtime's residency set keeps +// resident, so the kernel writes the engine's own memory instead of a CPU mirror. +struct Ds4EmbedArgsHost { + uint64_t rows; + uint64_t h; + uint32_t dim; + uint32_t hc; +}; + +bool embed_addr(const uint16_t* rows, int T, float* h, int dim, int hc) { + // the engine uploads emb_rows on ITS stream; this kernel runs on mtl::Context queue. Order the two before reading. + strata::gpu::stream_synchronize(nullptr); + if (std::getenv("DS4_DBG_ROWS")) { // what is the engine actually handing us as the embedding rows? + const uint16_t* rv = (const uint16_t*) strata::gpu::host_view((void*) rows); + std::fprintf(stderr, "DBG rows_in: ptr=%p view=%p first=%04x %04x %04x %04x\n", + (const void*) rows, (const void*) rv, + rv ? (unsigned) rv[0] : 0u, rv ? (unsigned) rv[1] : 0u, + rv ? (unsigned) rv[2] : 0u, rv ? (unsigned) rv[3] : 0u); + } + strata::dsv41::metal_host::shared_cx(); + const int lib = shared_lib(); + if (lib < 0) return false; + Ds4EmbedArgsHost a{}; + a.rows = (uint64_t) (uintptr_t) rows; + a.h = (uint64_t) (uintptr_t) h; + a.dim = (uint32_t) dim; + a.hc = (uint32_t) hc; + mtl::Pipeline pl; + try { pl = strata::dsv41::metal_host::shared_cx().pipeline(lib, "ds4_embed_j"); } catch (...) { return false; } + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::V(a)}; + d.grid[0] = (uint32_t) ((dim + 255) / 256); + d.grid[1] = (uint32_t) T; + d.grid[2] = 1; + d.tg[0] = 256; d.tg[1] = d.tg[2] = 1; + strata::dsv41::metal_host::shared_cx().run({d}); + strata::dsv41::metal_host::shared_cx().flush_runs(); + if (std::getenv("DS4_DBG_CMP")) { + const float* hv = (const float*) strata::gpu::host_view(h); + { + const size_t nv = (size_t) T * hc * dim; + size_t nf = 0, nz = 0; + for (size_t i = 0; i < nv; ++i) { const float v = hv[i]; if (!(v == v)) ++nf; else if (v == 0.0f) ++nz; } + std::fprintf(stderr, "DBG cmp: after embed n=%.0f nonfinite=%zu zeros=%zu\n", (double) nv, nf, nz); + } + std::fprintf(stderr, "DBG cmp: view %g %g %g | h=%g h+1=%g\n", + hv ? hv[0] : -1.0f, hv ? hv[1] : -1.0f, hv ? hv[2] : -1.0f, + hv ? hv[(size_t) dim * hc] : -1.0f, hv ? hv[(size_t) dim * hc + 1] : -1.0f); + } + if (std::getenv("DS4_DBG_ARGS")) { + mtl::Buffer bp = strata::dsv41::metal_host::shared_cx().wrap( + (void*) (((uintptr_t) strata::gpu::host_view(h)) & ~(uintptr_t) 16383), 16384); + mtl::Pipeline pp = strata::dsv41::metal_host::shared_cx().pipeline(lib, "ds4_args_probe"); + mtl::Dispatch pd; + pd.p = &pp; + pd.args = {mtl::Arg::B(bp, ((uintptr_t) strata::gpu::host_view(h)) & 16383), + mtl::Arg::V((unsigned) dim), mtl::Arg::V((unsigned) hc)}; + pd.grid[0] = 1; pd.grid[1] = pd.grid[2] = 1; pd.tg[0] = 32; pd.tg[1] = pd.tg[2] = 1; + strata::dsv41::metal_host::shared_cx().run({pd}); + strata::dsv41::metal_host::shared_cx().flush_runs(); + float got[3] = {0, 0, 0}; + { const float* hv2 = (const float*) strata::gpu::host_view(h); if (hv2) { got[0] = hv2[0]; got[1] = hv2[1]; } } + std::fprintf(stderr, "DBG args: kernel sees dim=%g hc=%g\n", got[0], got[1]); + } + return true; +} + +} // namespace strata::dsv41::metal_host \ No newline at end of file diff --git a/edge0-strata/metal/src/engine/exl3_legacy_stub.cpp b/edge0-strata/metal/src/engine/exl3_legacy_stub.cpp new file mode 100644 index 00000000..7d2cd4ce --- /dev/null +++ b/edge0-strata/metal/src/engine/exl3_legacy_stub.cpp @@ -0,0 +1,29 @@ +// metal/src/engine/exl3_legacy_stub.cpp - v1.1.8's EXL3 decode kernels, absent on Metal. +// +// exl3_legacy.cu exists for NVIDIA GPUs before sm_80, which pick it over the newer tiles; nothing on a Mac can select it +// (use_legacy() is false without CUDA), and the 2.25bpw text model does not use EXL3 at all. Leaving the file out of the +// target would break the link for these seven entry points, so they are defined here and stop the run naming themselves - +// the same contract as the stub generator, and a loud failure rather than a silent wrong decode. +#include +#include + +#include // MvJob, MoeArgs, WinArgs + +namespace strata::dsv41k::exl3::legacy { + +[[noreturn]] static void missing(const char* what) { + std::fprintf(stderr, + "strata metal: exl3::legacy::%s has no Metal implementation (v1.1.8's EXL3 kernels; the 2.25bpw text model " + "does not use EXL3)\n", what); + std::abort(); +} + +int launch_errors() { return 0; } +bool prepare(cudaStream_t) { missing("prepare"); } +bool several_ok() { return false; } +bool mv(const MvJob*, int, cudaStream_t) { missing("mv"); } +void moe_gate_up(const MoeArgs&, const float*, float*, cudaStream_t) { missing("moe_gate_up"); } +void moe_down(const MoeArgs&, const float*, const float*, float*, cudaStream_t) { missing("moe_down"); } +bool win_experts(const WinArgs&, cudaStream_t) { missing("win_experts"); } + +} // namespace strata::dsv41k::exl3::legacy diff --git a/edge0-strata/metal/src/engine/iq_host.cpp b/edge0-strata/metal/src/engine/iq_host.cpp new file mode 100644 index 00000000..28dca035 --- /dev/null +++ b/edge0-strata/metal/src/engine/iq_host.cpp @@ -0,0 +1,98 @@ +// metal/src/engine/iq_host.cpp - the host-side helpers of src/kernels/cuda/iq_kernels.cu the shared host code reads +// (format tables, row sizes, the native expert blob layout). The kernels themselves are not ported: on Metal the +// decode window runs mac-kern's chain (metal/src/engine/window.cpp). Copied from iq_kernels.cu; keep in step. +#include "strata/kernels/iq_kernels.hpp" + +#define GGML_COMMON_DECL_CPP +#include "ggml-common.h" +#include "ggml-common-tq.h" // block_tq2_t / block_tqk6 / block_tqk7 + +#include +#include + +namespace strata::kernels { +namespace { +// agentionai's trellis types (Gyro): TQ2_T (144), TQK6 (145), TQK7 (146); 128 weights a block +bool is_tq(int t) { return t == 144 || t == 145 || t == 146; } +bool is_iq(int t) { + return t == 16 || t == 17 || t == 18 || t == 20 || t == 21 || t == 22 || t == 23 || t == 29 || t == 42 || t == 11 || + t == 12 || t == 13 || t == 14 || t == 7 || t == 6 || t == 2 || t == 3 || t == 8; // 14: Q6_K (Unsloth's Q5/Q6 gate/up) +} +// values per block of the types the grouped expert kernels take (0 = none): STRATA_GU_FMTS / STRATA_D_FMTS +int gu_qk(int t) { + switch (t) { + case 16: case 17: case 18: case 21: case 22: case 23: case 29: case 12: case 13: case 14: return 256; + case 42: return 64; + case 144: case 145: return 128; // TQ2_T / TQK6 gate/up (the CUDA engine's STRATA_GU_FMTS) + case 6: case 2: case 3: case 8: return 32; + default: return 0; + } +} +int d_qk(int t) { + switch (t) { + case 23: return 256; + case 42: return 64; + case 144: case 146: return 128; // TQ2_T / TQK7 down + case 20: case 7: case 6: case 2: case 3: case 8: return 32; + default: return 0; + } +} +} // namespace + +bool iq_supported(int t) noexcept { return is_iq(t); } +bool embed_type_supported(int t) noexcept { return is_iq(t) || t == 30; } + +size_t iq_row_bytes(int t, int64_t n) noexcept { + switch (t) { + case 16: return (size_t) (n / 256) * sizeof(block_iq2_xxs); + case 17: return (size_t) (n / 256) * sizeof(block_iq2_xs); + case 18: return (size_t) (n / 256) * sizeof(block_iq3_xxs); + case 20: return (size_t) (n / 32) * sizeof(block_iq4_nl); + case 21: return (size_t) (n / 256) * sizeof(block_iq3_s); + case 22: return (size_t) (n / 256) * sizeof(block_iq2_s); + case 29: return (size_t) (n / 256) * sizeof(block_iq1_m); + case 23: return (size_t) (n / 256) * sizeof(block_iq4_xs); + case 11: return (size_t) (n / 256) * sizeof(block_q3_K); + case 42: return (size_t) (n / 64) * sizeof(block_q2_0); + case 12: return (size_t) (n / 256) * sizeof(block_q4_K); + case 13: return (size_t) (n / 256) * sizeof(block_q5_K); + case 14: return (size_t) (n / 256) * sizeof(block_q6_K); + case 7: return (size_t) (n / 32) * sizeof(block_q5_1); + case 6: return (size_t) (n / 32) * sizeof(block_q5_0); + case 2: return (size_t) (n / 32) * sizeof(block_q4_0); + case 3: return (size_t) (n / 32) * sizeof(block_q4_1); + case 8: return (size_t) (n / 32) * sizeof(block_q8_0); + case 144: return (size_t) (n / 128) * sizeof(block_tq2_t); + case 145: return (size_t) (n / 128) * sizeof(block_tqk6); + case 146: return (size_t) (n / 128) * sizeof(block_tqk7); + case 30: return (size_t) n * 2; + default: return 0; + } +} + +bool native_expert_supported(int gu_type, int d_type, int64_t n_embd, int64_t n_ff) noexcept { + const int qg = gu_qk(gu_type), qd = d_qk(d_type); + return qg > 0 && qd > 0 && (is_iq(gu_type) || is_tq(gu_type)) && (is_iq(d_type) || is_tq(d_type)) && n_embd % qg == 0 && n_ff % qd == 0 && + n_embd % 256 == 0 && (n_ff * n_embd) % 256 == 0; // (the trellis rows: 128-value blocks, 2560 = 20, 640 = 5) +} + +NativeExpertLayout native_expert_layout(int gu_type, int d_type, int64_t n_embd, int64_t n_ff) { + NativeExpertLayout L; + L.gu_type = gu_type; + L.d_type = d_type; + L.n_embd = n_embd; + L.n_ff = n_ff; + L.gu_row = iq_row_bytes(gu_type, n_embd); + L.d_row = iq_row_bytes(d_type, n_ff); + L.up_off = (size_t) n_ff * L.gu_row; + L.down_off = 2 * L.up_off; + L.bytes = L.down_off + (size_t) n_embd * L.d_row; + return L; +} + +size_t native_expert_scratch_bytes(int64_t cap, int64_t n_ff) { + const size_t f = (size_t) cap * (size_t) n_ff * sizeof(float); + return 3 * ((f + 255) & ~(size_t) 255) + (((size_t) cap * (size_t) (n_ff / 32) * sizeof(block_q8_1) + 255) & ~(size_t) 255); +} + +} // namespace strata::kernels diff --git a/edge0-strata/metal/src/engine/latent_stubs.cpp b/edge0-strata/metal/src/engine/latent_stubs.cpp new file mode 100644 index 00000000..6d300e21 --- /dev/null +++ b/edge0-strata/metal/src/engine/latent_stubs.cpp @@ -0,0 +1,152 @@ +// Shared host code links these entry points. The Metal window owns latent execution. +#include "strata/kernels/elementwise.hpp" +#include "strata/kernels/iq_kernels.hpp" +#include "strata/kernels/latent.hpp" +#include "strata/kernels/s2_expert_grouped.hpp" +#include "strata/kernels/cpu/kq_avx2.hpp" +#include "strata/prefill/moe_fused.hpp" +#include "strata/prefill/moe_fused_q3k.hpp" +#include "strata/prefill/lat_q3k_fused.hpp" + +namespace strata::metal_engine { +void missing(const char* name); +} + +namespace strata::kernels { + +void copy_bytes_from_mapped(void*, const void*, int64_t, void*) { + strata::metal_engine::missing("strata::kernels::copy_bytes_from_mapped"); +} + +void doorbell_publish_value(const float*, const int32_t*, const float*, int64_t, int64_t, float*, int32_t*, + float*, uint32_t*, uint32_t, void*) { + strata::metal_engine::missing("strata::kernels::doorbell_publish_value"); +} + +void iq_mmvq(int, const void*, const void*, float*, int, int, int, void*) { + strata::metal_engine::missing("strata::kernels::iq_mmvq"); +} + +void lat_colscale_rows(float*, int64_t, int, const float*, void*) { + strata::metal_engine::missing("strata::kernels::lat_colscale_rows"); +} + +const float* lat_gu_colscale_dev(int, int) { + strata::metal_engine::missing("strata::kernels::lat_gu_colscale_dev"); + return {}; +} + +void latent_plan_rows(int32_t*, const int32_t*, const int32_t*, const int32_t*, const int32_t*, int64_t, int, + int, int, void*) { + strata::metal_engine::missing("strata::kernels::latent_plan_rows"); +} + +void latent_u_grouped(const uint8_t*, const uint8_t*, const int32_t*, const int32_t*, int64_t, int64_t, int, + int, float*, void*, int) { + strata::metal_engine::missing("strata::kernels::latent_u_grouped"); +} + +void latent_u_matvec(const uint8_t*, const uint8_t*, const int32_t*, const int32_t*, const int32_t*, + const int32_t*, int64_t, int64_t, int, int, int, float*, void*, int) { + strata::metal_engine::missing("strata::kernels::latent_u_matvec"); +} + +void moe_hit_add_mapped(float*, const float*, const volatile float*, const int32_t*, const int32_t*, int64_t, + int64_t, const volatile uint32_t*, uint32_t, void*) { + strata::metal_engine::missing("strata::kernels::moe_hit_add_mapped"); +} + +void moe_hit_set(float*, const float*, const int32_t*, const int32_t*, int64_t, int64_t, void*) { + strata::metal_engine::missing("strata::kernels::moe_hit_set"); +} + +void quantize_q8_0_rows(const float*, long long, long long, void*, void*) { + strata::metal_engine::missing("strata::kernels::quantize_q8_0_rows"); +} + +} // namespace strata::kernels + +namespace strata::kernels::cpu { + +void lat_q4_0_q8_1_rows(const uint8_t*, size_t, int, const void*const*, int, float*const*, int, int) { + strata::metal_engine::missing("strata::kernels::cpu::lat_q4_0_q8_1_rows"); +} + +void lat_q8_0_q8_1_rows(const uint8_t*, size_t, int, const void*const*, int, float*const*, int, int) { + strata::metal_engine::missing("strata::kernels::cpu::lat_q8_0_q8_1_rows"); +} + +void lat_quant_q8_1_qs_avx2(const float*, const float*, int8_t*) { + strata::metal_engine::missing("strata::kernels::cpu::lat_quant_q8_1_qs_avx2"); +} + +} // namespace strata::kernels::cpu + +namespace strata::prefill::fused { + +void experts_latent(const WideBatch&, const LatentGeom&, int, int64_t, const void*, const void*, + const int32_t*, void*, float*, void*) { + strata::metal_engine::missing("strata::prefill::fused::experts_latent"); +} + +void group_latent(const int32_t*, const int32_t*, int, void*, size_t) { + strata::metal_engine::missing("strata::prefill::fused::group_latent"); +} + +void latent_dense_d(const void*, const void*, int64_t, int, int, int64_t, void*, void*, int) { + strata::metal_engine::missing("strata::prefill::fused::latent_dense_d"); +} + +void latent_dense_u(const void*, const void*, const int32_t*, const int32_t*, int, int64_t, int64_t, int, + int, float*, void*, int) { + strata::metal_engine::missing("strata::prefill::fused::latent_dense_u"); +} + +void latent_dense_u_tables(const int32_t*, const int32_t*, int64_t, int, int32_t*, int32_t*, int32_t*, + void*) { + strata::metal_engine::missing("strata::prefill::fused::latent_dense_u_tables"); +} + +bool latent_fused_available() { + strata::metal_engine::missing("strata::prefill::fused::latent_fused_available"); + return {}; +} + +size_t latent_group_bytes(int64_t, int) { + strata::metal_engine::missing("strata::prefill::fused::latent_group_bytes"); + return {}; +} + +size_t latent_h_off(int64_t, int) { + strata::metal_engine::missing("strata::prefill::fused::latent_h_off"); + return {}; +} + +size_t latent_scratch_bytes(int64_t, int, int64_t) { + strata::metal_engine::missing("strata::prefill::fused::latent_scratch_bytes"); + return {}; +} + +} // namespace strata::prefill::fused + +namespace strata::prefill::fused_q3k { + +void experts_latent_q3k(const fused::WideBatch&, const fused::LatentGeom&, int, int64_t, const void*, + const void*, const int32_t*, void*, float*, void*) { + strata::metal_engine::missing("strata::prefill::fused_q3k::experts_latent_q3k"); +} + +void latent_dense_d(const void*, const void*, int64_t, int, int, int64_t, void*, void*, int, const float*) { + strata::metal_engine::missing("strata::prefill::fused_q3k::latent_dense_d"); +} + +} // namespace strata::prefill::fused_q3k + +namespace strata::prefill { + +void lat_q3k_fused(const uint8_t*const*, const int32_t*, int, const void*, void*, float*, int64_t, void*, + const int32_t*, const int32_t*, const int32_t*) { + strata::metal_engine::missing("strata::prefill::lat_q3k_fused"); +} + +} // namespace strata::prefill diff --git a/edge0-strata/metal/src/engine/missing.cpp b/edge0-strata/metal/src/engine/missing.cpp new file mode 100644 index 00000000..27ecec2d --- /dev/null +++ b/edge0-strata/metal/src/engine/missing.cpp @@ -0,0 +1,80 @@ +// metal/src/engine/missing.cpp - what a kernel entry point that has no Metal implementation yet does when called. +// +// Default: say which one and stop (exit code 86), so a run never continues on outputs nobody wrote. +// STRATA_METAL_STUBS=log: say it once per entry point and return (outputs untouched) - for listing which entry points a +// code path reaches, never for answers. +#include +#include +#include +#include +#include + +namespace strata::metal_engine { + +void missing(const char* name) { + static const bool log_only = [] { + const char* v = std::getenv("STRATA_METAL_STUBS"); + return v != nullptr && std::string(v) == "log"; + }(); + if (!log_only) { + std::fprintf(stderr, "strata metal: %s has no Metal implementation yet\n", name); + std::fflush(stderr); + std::_Exit(86); + } + static std::mutex mu; + static std::set seen; + std::lock_guard lk(mu); + if (seen.insert(name).second) std::fprintf(stderr, "strata metal: STUB %s\n", name); +} + +} // namespace strata::metal_engine + +// ---- entry points the shared host loop asks every request, with nothing to do on Metal yet +#include "strata/kernels/cvec.hpp" +namespace strata::kernels { +// no control vector on Metal (--control-vector-scaled is not ported): the empty one, `loaded()` false +const Cvec& cvec() { + static const Cvec none{}; + return none; +} +} // namespace strata::kernels + +// an uncaught exception or abort on the Metal build prints where it came from (macOS gives no core dump by default and +// lldb needs a debugging permission a user's session may not have) +#include +#include +#include +#include +#include +namespace { +[[noreturn]] void metal_terminate() { + void* fr[64]; + const int n = backtrace(fr, 64); + const char msg[] = "strata metal: terminate - the call stack:\n"; + (void) !write(2, msg, sizeof msg - 1); + backtrace_symbols_fd(fr, n, 2); + std::abort(); +} +const bool g_terminate_set = (std::set_terminate(metal_terminate), true); +// a crash (SIGSEGV / SIGBUS) prints its call stack too: macOS keeps no report for a command-line process by default +void metal_crash(int sig, siginfo_t* si, void*) { + void* fr[64]; + const int n = backtrace(fr, 64); + char msg[160]; + const int m = std::snprintf(msg, sizeof msg, "strata metal: crashed (signal %d at address %p) - the call stack:\n", sig, + si ? si->si_addr : nullptr); + (void) !write(2, msg, (size_t) m); + backtrace_symbols_fd(fr, n, 2); + signal(sig, SIG_DFL); + raise(sig); +} +bool set_crash() { + struct sigaction sa {}; + sa.sa_sigaction = metal_crash; + sa.sa_flags = SA_SIGINFO; + sigaction(SIGSEGV, &sa, nullptr); + sigaction(SIGBUS, &sa, nullptr); + return true; +} +const bool g_crash_set = set_crash(); +} // namespace diff --git a/edge0-strata/metal/src/engine/mmvq_host.cpp b/edge0-strata/metal/src/engine/mmvq_host.cpp new file mode 100644 index 00000000..af39859b --- /dev/null +++ b/edge0-strata/metal/src/engine/mmvq_host.cpp @@ -0,0 +1,66 @@ +// metal/src/engine/mmvq_host.cpp - the host-side helpers of src/kernels/cuda/native_mmvq.cu the shared host code reads +// (which GGUF types the dense projections take, their byte sizes). Copied from native_mmvq.cu; keep in step. The +// kernels are not ported: on Metal the decode window runs mac-kern's GEMVs (metal/src/engine/window.cpp). +#include "strata/kernels/iq_kernels.hpp" +#include "strata/kernels/native_mmvq.hpp" + +#include +#include +#include +#include + +namespace strata::kernels { +namespace { +constexpr int Q8K = 32; +constexpr std::size_t kQ81Block = 4 + 32; // half2 ds + int8 qs[32] +void validate_shape(int n_in, int ncols, int block_elems = Q8K) { + if (n_in <= 0 || n_in % block_elems != 0) + throw std::invalid_argument("native MMVQ requires n_in > 0 and divisible by its block element count"); + if (ncols < 1 || ncols > 8) throw std::invalid_argument("native MMVQ requires 1 <= ncols <= 8"); +} +} // namespace + +bool native_mmvq_supported(int ggml_type) noexcept { + return ggml_type == 2 || ggml_type == 6 || ggml_type == 7 || ggml_type == 8 || ggml_type == 11 || + ggml_type == 12 || ggml_type == 13 || ggml_type == 14 || ggml_type == 20 || + ggml_type == 23 || ggml_type == 42 || ggml_type == 16 || ggml_type == 17 || ggml_type == 18 || + ggml_type == 21 || ggml_type == 22 || ggml_type == 29; +} + +std::size_t native_mmvq_weight_bytes(int ggml_type, int n_in, int n_out) { + int block_elems, block_bytes; + switch (ggml_type) { + case 2: block_elems = 32; block_bytes = 18; break; + case 6: block_elems = 32; block_bytes = 22; break; + case 7: block_elems = 32; block_bytes = 24; break; + case 8: block_elems = 32; block_bytes = 34; break; + case 20: block_elems = 32; block_bytes = 18; break; + case 11: block_elems = 256; block_bytes = 110; break; + case 12: block_elems = 256; block_bytes = 144; break; + case 13: block_elems = 256; block_bytes = 176; break; + case 14: block_elems = 256; block_bytes = 210; break; + case 23: block_elems = 256; block_bytes = 136; break; + case 42: block_elems = 64; block_bytes = 18; break; + case 16: case 17: case 18: case 21: case 22: case 29: + block_elems = 256; block_bytes = (int) iq_row_bytes(ggml_type, 256); break; + default: throw std::invalid_argument("unsupported native MMVQ GGML type"); + } + validate_shape(n_in, 1, block_elems); + if (n_out <= 0) throw std::invalid_argument("native MMVQ requires n_out > 0"); + const std::size_t row_bytes = std::size_t(n_in / block_elems) * block_bytes; + if (row_bytes > std::numeric_limits::max() / std::size_t(n_out)) + throw std::length_error("native MMVQ weight byte count overflows size_t"); + return row_bytes * std::size_t(n_out); +} + +std::size_t native_q8_1_bytes(int n_in, int ncols) { + validate_shape(n_in, ncols); + return std::size_t(ncols) * std::size_t(n_in / Q8K) * kQ81Block; +} +std::size_t native_q8_1_il_bytes(int n_in, int ncols) { return native_q8_1_bytes(n_in, ncols) * 2; } +bool native_mmvq_il_supported(int, int, int) { return false; } +bool native_q8_0_packed_eligible(int, int) { return false; } +bool native_q8_0_packed_enabled() { return false; } +bool native_q6_k_packed_enabled() { return false; } + +} // namespace strata::kernels diff --git a/edge0-strata/metal/src/engine/msl/ds4_init.metal b/edge0-strata/metal/src/engine/msl/ds4_init.metal new file mode 100644 index 00000000..87ec57a6 --- /dev/null +++ b/edge0-strata/metal/src/engine/msl/ds4_init.metal @@ -0,0 +1,2800 @@ +// metal/src/engine/msl/ds4_init.metal - the first kernels the engine asks for, ported from ds4_kernels.cu. +// +// ds4_kernels.cu:1477 __global__ void init_pre_kernel(float* pre, int T) { +// const int i = blockIdx.x * blockDim.x + threadIdx.x; +// if (i < T * 8) pre[i] = (i % 8) == 0 ? 1.0f : 0.0f; +// } +#include + +using namespace metal; + +// the grid is threadgroups of 128; the flat id is what CUDA's blockIdx*blockDim+threadIdx was +kernel void ds4_init_pre(device float* pre [[buffer(0)]], constant uint& T [[buffer(1)]], + uint gid [[thread_position_in_grid]]) { + if (gid < T * 8u) pre[gid] = (gid % 8u) == 0u ? 1.0f : 0.0f; +} + +// ds4_kernels.cu:130 embed_kernel: v = bf16(rows[t*dim + e]); h[(t*hc + s)*dim + e] = v for each of hc copies +kernel void ds4_embed(device const ushort* rows [[buffer(0)]], device float* h [[buffer(1)]], + constant uint& dim [[buffer(2)]], constant uint& hc [[buffer(3)]], + uint2 tgp [[threadgroup_position_in_grid]], ushort lane [[thread_index_in_threadgroup]]) { + const uint e = tgp.x * 256u + uint(lane); + if (e >= dim) return; + const uint t = tgp.y; + const uint bits = uint(rows[(size_t) t * dim + e]) << 16; // BF16 -> F32 is the top half + const float v = as_type(bits); + for (uint s = 0; s < hc; ++s) h[((size_t) t * hc + s) * dim + e] = v; +} + +// ds4_model.cu:2737 cublasSgemm(OP_T, OP_N, n_mix, T, hcd, ...) - the mHC mixes: mixes[t][m] = h[t] . fn[m] +// one thread an output element; both operands F32, alpha = 1, beta = 0 (every call site seen) +kernel void ds4_sgemm_tn(device const float* A [[buffer(0)]], device const float* B [[buffer(1)]], + device float* C [[buffer(2)]], constant uint& m [[buffer(3)]], constant uint& nt [[buffer(4)]], + constant uint& k [[buffer(5)]], constant uint& lda [[buffer(6)]], constant uint& ldb [[buffer(7)]], + constant uint& ldc [[buffer(8)]], uint gid [[thread_position_in_grid]]) { + const uint total = m * nt; + if (gid >= total) return; + const uint t = gid / m, mm = gid % m; + const device float* a = A + (size_t) mm * lda; // row mm of A + const device float* b = B + (size_t) t * ldb; // row t of B + float acc = 0.0f; + for (uint i = 0; i < k; ++i) acc += a[i] * b[i]; + C[(size_t) t * ldc + mm] = acc; +} + +// ds4_kernels.cu:1894 finish_kernel - RMS norm (f.w), mrope (f.rd), rounding (f.quant), then store to x or dst. +// The mrope and rounding branches are not written yet; the host wrapper will decline until they are. +struct Ds4Finish { + device float* w; + float eps; + int rd, pos, inverse, vpp, pos_step; + device float* freqs; + int quant; + device ushort* dst; + long dst_stride; + device uchar* pdst; + long pdst_stride; +}; + +static inline float ds4_warp_sum(float v) { + v += simd_shuffle_xor(v, 16); + v += simd_shuffle_xor(v, 8); + v += simd_shuffle_xor(v, 4); + v += simd_shuffle_xor(v, 2); + v += simd_shuffle_xor(v, 1); + return v; +} +static inline float ds4_block_sum(float v, threadgroup float* red, ushort lane, ushort w) { + v = ds4_warp_sum(v); + if (lane == 0) red[w] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (w == 0) { + float t = lane < 8 ? red[lane] : 0.0f; + t = ds4_warp_sum(t); + if (lane == 0) red[0] = t; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + return red[0]; +} + +// ds4_kernels.cu:65 pow2_ceil / :88 fp8_rt / :90 e2m1_rt - the KV quantiser's scale and roundings, shared by finish_pack +// (the packed rows) and ds4_finish's rounding branch +static inline float ds4_pow2_ceil(float x) { + const uint b = as_type(x); + const int e = int((b >> 23) & 0xffu) - 127 + ((b & 0x7fffffu) != 0u ? 1 : 0); + return as_type(uint(e + 127) << 23); +} + +static inline float ds4_e4m3_round(float v) { // E4M3, nearest even (the input is clamped) + const float a = fabs(v); + if (!(a == a)) return v; + float r; + if (a >= 448.0f) r = 448.0f; + else { + const int e = int((as_type(a) >> 23) & 0xffu) - 127; + const float step = e < -6 ? 0.001953125f : as_type(uint(e - 3 + 127) << 23); + r = a == 0.0f ? 0.0f : rint(a / step) * step; + r = fmin(r, 448.0f); + } + return copysign(r, v); +} + +static inline float ds4_e2m1_rt(float v) { // E2M1 magnitudes 0 .5 1 1.5 2 3 4 6, nearest, ties to even + const float a = fabs(v); + float r; + if (a <= 0.25f) r = 0.0f; + else if (a < 0.75f) r = 0.5f; + else if (a <= 1.25f) r = 1.0f; + else if (a < 1.75f) r = 1.5f; + else if (a <= 2.5f) r = 2.0f; + else if (a < 3.5f) r = 3.0f; + else if (a <= 5.0f) r = 4.0f; + else r = 6.0f; + return copysign(r, v); +} + +kernel void ds4_finish(device float* x [[buffer(0)]], constant long& stride [[buffer(1)]], + constant uint& len [[buffer(2)]], constant Ds4Finish& f [[buffer(3)]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], + ushort w [[simdgroup_index_in_threadgroup]], uint b [[threadgroup_position_in_grid]]) { + threadgroup float red[8]; + device float* v = x + (long) b * stride; + const uint nj = (len + 255u) / 256u; + float r[32]; + float ss = 0.0f; + for (uint j = 0; j < 32u; ++j) { + if (j >= nj) break; + const uint e = tid + j * 256u; + r[j] = e < len ? v[e] : 0.0f; + ss += r[j] * r[j]; + } + if (f.w != nullptr) { + const float tot = ds4_block_sum(ss, red, lane, w); + const float inv = rsqrt(tot / float(len) + f.eps); + for (uint j = 0; j < 32u; ++j) { + if (j >= nj) break; + const uint e = tid + j * 256u; + if (e < len) r[j] = r[j] * inv * f.w[e]; + } + } + if (f.rd > 0) { + // mrope: the last rd values move in (re, im) pairs a lane apart (ds4_kernels.cu:161): the partner comes from a + // lane xor 1, the angle from the token's position and the frequency of the pair + const uint r0 = len - uint(f.rd); + const uint pos = uint(f.pos) + (b / uint(f.vpp)) * uint(f.pos_step); + for (uint j = 0; j < 32u; ++j) { + if (j >= nj) break; + const uint e = tid + j * 256u; + const float other = simd_shuffle_xor(r[j], 1); + if (e >= r0 && e < len) { + const uint i = (e - r0) >> 1; + const float ang = float(pos) * f.freqs[i]; + const float c = cos(ang); + const float sn = f.inverse ? -sin(ang) : sin(ang); + r[j] = ((e - r0) & 1u) == 0u ? r[j] * c - other * sn : other * sn + r[j] * c; + } + } + } + if (f.quant != 0) { // the KV rounding without the pack: the group's scale, then the value rounded to the format + const uint g = f.quant == 2 ? 16u : 32u; + for (uint j = 0; j < 32u; ++j) { + if (j >= nj) break; + const uint e = tid + j * 256u; + float a = e < len ? fabs(r[j]) : 0.0f; + for (uint o = g / 2u; o > 0u; o >>= 1) a = fmax(a, simd_shuffle_xor(a, ushort(o))); + if (e >= len) continue; + if (f.quant == 1) { + const float s = ds4_pow2_ceil(fmax(a, 1e-4f) * (1.0f / 448.0f)); + r[j] = ds4_e4m3_round(fmin(fmax(r[j] / s, -448.0f), 448.0f)) * s; + } else if (f.quant == 2) { + const float s = ds4_e4m3_round(fmax(a, 6.0f * 0.001953125f) / 6.0f); + r[j] = ds4_e2m1_rt(fmin(fmax(r[j] / s, -6.0f), 6.0f)) * s; + } else { + const float s = ds4_pow2_ceil(fmax(a, 6.0f * 1.1754943508222875e-38f) * (1.0f / 6.0f)); + r[j] = ds4_e2m1_rt(fmin(fmax(r[j] / s, -6.0f), 6.0f)) * s; + } + } + } + for (uint j = 0; j < 32u; ++j) { + if (j >= nj) break; + const uint e = tid + j * 256u; + if (e >= len) continue; + if (f.dst != nullptr) { + const uint bits = as_type(r[j]); + f.dst[(long) b * f.dst_stride + (long) e] = ushort((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); // BF16 round to nearest even, the vector's own row + } else { + v[e] = r[j]; + } + } +} + +// ds4_kernels.cu:1677 hc_post_inplace_kernel: h[t][k][e] = post[t][k]*y[t][e] + sum_j comb[t][j][k]*h_old[t][j][e] +kernel void ds4_hc_post(device const float* y [[buffer(0)]], device float* h [[buffer(1)]], + device const float* post [[buffer(2)]], device const float* comb [[buffer(3)]], + constant uint& dim [[buffer(4)]], constant uint& hc [[buffer(5)]], + uint2 tgp [[threadgroup_position_in_grid]], ushort lane [[thread_index_in_threadgroup]]) { + const uint e = tgp.x * 256u + uint(lane); + if (e >= dim) return; + const uint t = tgp.y; + device float* r = h + (size_t) t * hc * dim; + float rv[8]; + for (uint j = 0; j < 8u; ++j) if (j < hc) rv[j] = r[(size_t) j * dim + e]; + const float yv = y[(size_t) t * dim + e]; + for (uint k = 0; k < 8u; ++k) { + if (k >= hc) break; + float a = post[(size_t) t * 8u + k] * yv; + for (uint j = 0; j < 8u; ++j) if (j < hc) a += comb[(size_t) t * 64u + j * hc + k] * rv[j]; + r[(size_t) k * dim + e] = a; + } +} + +// ds4_kernels.cu:1692 hc_mix_scale_kernel: mixes[t][i] *= rsqrt(sum(h[t])^2 / hcdim + eps), one threadgroup a token +kernel void ds4_hc_mix_scale(device const float* h [[buffer(0)]], device float* mixes [[buffer(1)]], + constant uint& hcdim [[buffer(2)]], constant uint& n_mix [[buffer(3)]], + constant float& eps [[buffer(4)]], uint b [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], + ushort w [[simdgroup_index_in_threadgroup]]) { + threadgroup float red[8]; + const device float* x = h + (size_t) b * hcdim; + float ss = 0.0f; + for (uint k = tid * 4u; k < hcdim; k += 256u * 4u) { + const float4 hv = *(const device float4*) (x + k); + ss += hv.x * hv.x + hv.y * hv.y + hv.z * hv.z + hv.w * hv.w; + } + ss = ds4_block_sum(ss, red, lane, w); + const float inv = rsqrt(ss / float(hcdim) + eps); + if (tid < n_mix) mixes[(size_t) b * 32u + tid] *= inv; +} + +// ds4_kernels.cu:1017 hc_mix_kernel: mixes[t][m] = (h[t] . fn[m]) * rsqrt(mean(h[t]^2) + eps), one threadgroup a (m, t) +kernel void ds4_hc_mix(device const float* h [[buffer(0)]], device const float* fn [[buffer(1)]], + device float* mixes [[buffer(2)]], constant uint& hcdim [[buffer(3)]], + constant float& eps [[buffer(4)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], + ushort w [[simdgroup_index_in_threadgroup]]) { + threadgroup float red[8]; + const uint m = tgp.x, t = tgp.y; + const device float* x = h + (size_t) t * hcdim; + const device float* f = fn + (size_t) m * hcdim; + float acc = 0.0f, ss = 0.0f; + for (uint k = tid * 4u; k < hcdim; k += 256u * 4u) { + const float4 hv = *(const device float4*) (x + k); + const float4 fv = *(const device float4*) (f + k); + acc += hv.x * fv.x + hv.y * fv.y + hv.z * fv.z + hv.w * fv.w; + ss += hv.x * hv.x + hv.y * hv.y + hv.z * hv.z + hv.w * hv.w; + } + acc = ds4_block_sum(acc, red, lane, w); + ss = ds4_block_sum(ss, red, lane, w); + if (tid == 0) mixes[(size_t) t * 32u + m] = acc * rsqrt(ss / float(hcdim) + eps); +} + +// hc_mix for a prompt chunk: one threadgroup a token for four mixes a pass (the token's row is read once per four mixes instead of once per +// mix). Every mix's partial sums keep the single-mix kernel's order, so the result is bit-identical. +kernel void ds4_hc_mix_all(device const float* h [[buffer(0)]], device const float* fn [[buffer(1)]], + device float* mixes [[buffer(2)]], constant uint& hcdim [[buffer(3)]], + constant float& eps [[buffer(4)]], constant uint& n_mix [[buffer(5)]], + uint t [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], + ushort w [[simdgroup_index_in_threadgroup]]) { + threadgroup float red[8]; + const device float* x = h + (size_t) t * hcdim; + float ss = 0.0f; + for (uint k = tid * 4u; k < hcdim; k += 256u * 4u) { + const float4 hv = *(const device float4*) (x + k); + ss += hv.x * hv.x + hv.y * hv.y + hv.z * hv.z + hv.w * hv.w; + } + ss = ds4_block_sum(ss, red, lane, w); + const float inv = rsqrt(ss / float(hcdim) + eps); + for (uint m0 = 0; m0 < n_mix; m0 += 4u) { + const uint nm = min(4u, n_mix - m0); + float a0 = 0.0f, a1 = 0.0f, a2 = 0.0f, a3 = 0.0f; + const device float* f0 = fn + (size_t) m0 * hcdim; + const device float* f1 = f0 + (nm > 1u ? hcdim : 0u); + const device float* f2 = f0 + (nm > 2u ? 2u * hcdim : 0u); + const device float* f3 = f0 + (nm > 3u ? 3u * hcdim : 0u); + for (uint k = tid * 4u; k < hcdim; k += 256u * 4u) { + const float4 hv = *(const device float4*) (x + k); + const float4 v0 = *(const device float4*) (f0 + k); + const float4 v1 = *(const device float4*) (f1 + k); + const float4 v2 = *(const device float4*) (f2 + k); + const float4 v3 = *(const device float4*) (f3 + k); + a0 += hv.x * v0.x + hv.y * v0.y + hv.z * v0.z + hv.w * v0.w; + a1 += hv.x * v1.x + hv.y * v1.y + hv.z * v1.z + hv.w * v1.w; + a2 += hv.x * v2.x + hv.y * v2.y + hv.z * v2.z + hv.w * v2.w; + a3 += hv.x * v3.x + hv.y * v3.y + hv.z * v3.z + hv.w * v3.w; + } + a0 = ds4_block_sum(a0, red, lane, w); + a1 = ds4_block_sum(a1, red, lane, w); + a2 = ds4_block_sum(a2, red, lane, w); + a3 = ds4_block_sum(a3, red, lane, w); + if (tid == 0) { + mixes[(size_t) t * 32u + m0] = a0 * inv; + if (nm > 1u) mixes[(size_t) t * 32u + m0 + 1u] = a1 * inv; + if (nm > 2u) mixes[(size_t) t * 32u + m0 + 2u] = a2 * inv; + if (nm > 3u) mixes[(size_t) t * 32u + m0 + 3u] = a3 * inv; + } + } +} + +// hc mix + split in one dispatch. The mix is ds4_hc_mix_all's (four mixes a pass, every one keeping the single-mix +// kernel's summation order, so the values are bit-identical), then the same threadgroup runs the sinkhorn split of the +// 24 mixes into pre / post / comb. Two dispatches become one, and nothing else changes. +kernel void ds4_hc_mix_split(device const float* h [[buffer(0)]], device const float* fn [[buffer(1)]], + device float* mixes [[buffer(2)]], constant uint& hcdim [[buffer(3)]], + constant float& eps [[buffer(4)]], constant uint& n_mix [[buffer(5)]], + device const float* scale [[buffer(6)]], device const float* base [[buffer(7)]], + device float* pre [[buffer(8)]], device float* post [[buffer(9)]], + device float* comb [[buffer(10)]], constant uint& hc [[buffer(11)]], + constant uint& iters [[buffer(12)]], constant float& hc_eps [[buffer(13)]], + uint t [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort wid [[simdgroup_index_in_threadgroup]]) { + threadgroup float red[8]; + const device float* x = h + (size_t) t * hcdim; + float ss = 0.0f; + for (uint k = tid * 4u; k < hcdim; k += 256u * 4u) { + const float4 hv = *(const device float4*) (x + k); + ss += hv.x * hv.x + hv.y * hv.y + hv.z * hv.z + hv.w * hv.w; + } + ss = ds4_block_sum(ss, red, lane, wid); + const float inv = rsqrt(ss / float(hcdim) + eps); + for (uint m0 = 0; m0 < n_mix; m0 += 4u) { + const uint nm = min(4u, n_mix - m0); + float a0 = 0.0f, a1 = 0.0f, a2 = 0.0f, a3 = 0.0f; + const device float* f0 = fn + (size_t) m0 * hcdim; + const device float* f1 = f0 + (nm > 1u ? hcdim : 0u); + const device float* f2 = f0 + (nm > 2u ? 2u * hcdim : 0u); + const device float* f3 = f0 + (nm > 3u ? 3u * hcdim : 0u); + for (uint k = tid * 4u; k < hcdim; k += 256u * 4u) { + const float4 hv = *(const device float4*) (x + k); + const float4 v0 = *(const device float4*) (f0 + k); + const float4 v1 = *(const device float4*) (f1 + k); + const float4 v2 = *(const device float4*) (f2 + k); + const float4 v3 = *(const device float4*) (f3 + k); + a0 += hv.x * v0.x + hv.y * v0.y + hv.z * v0.z + hv.w * v0.w; + a1 += hv.x * v1.x + hv.y * v1.y + hv.z * v1.z + hv.w * v1.w; + a2 += hv.x * v2.x + hv.y * v2.y + hv.z * v2.z + hv.w * v2.w; + a3 += hv.x * v3.x + hv.y * v3.y + hv.z * v3.z + hv.w * v3.w; + } + a0 = ds4_block_sum(a0, red, lane, wid); + a1 = ds4_block_sum(a1, red, lane, wid); + a2 = ds4_block_sum(a2, red, lane, wid); + a3 = ds4_block_sum(a3, red, lane, wid); + if (tid == 0) { + mixes[(size_t) t * 32u + m0] = a0 * inv; + if (nm > 1u) mixes[(size_t) t * 32u + m0 + 1u] = a1 * inv; + if (nm > 2u) mixes[(size_t) t * 32u + m0 + 2u] = a2 * inv; + if (nm > 3u) mixes[(size_t) t * 32u + m0 + 3u] = a3 * inv; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + // the split: ds4_hc_split4's math, in the first warp (every lane of it takes the same path) + const device float* mx = mixes + (size_t) t * 32u; + if (tid < 4u) { + pre[(size_t) t * 8u + tid] = 1.0f / (1.0f + exp(-(mx[tid] * scale[0] + base[tid]))) + hc_eps; + post[(size_t) t * 8u + tid] = 2.0f / (1.0f + exp(-(mx[4u + tid] * scale[1] + base[4u + tid]))); + } + if (tid < 32u) { + const bool on = lane < 16u; + float c = on ? mx[8u + lane] * scale[2] + base[8u + lane] : -3.402823466e+38f; + float m = fmax(c, simd_shuffle_xor(c, ushort(1))); + m = fmax(m, simd_shuffle_xor(m, ushort(2))); + c = on ? exp(c - m) : 0.0f; + float rs = c + simd_shuffle_xor(c, ushort(1)); + rs += simd_shuffle_xor(rs, ushort(2)); + c = c / rs + hc_eps; + c = on ? c : 0.0f; + float cs = c + simd_shuffle_xor(c, ushort(4)); + cs += simd_shuffle_xor(cs, ushort(8)); + c = c / (cs + hc_eps); + for (uint it = 1; it < iters; ++it) { + float r = c + simd_shuffle_xor(c, ushort(1)); + r += simd_shuffle_xor(r, ushort(2)); + c = c / (r + hc_eps); + float q = c + simd_shuffle_xor(c, ushort(4)); + q += simd_shuffle_xor(q, ushort(8)); + c = c / (q + hc_eps); + } + if (on) comb[(size_t) t * 64u + lane] = c; + } +} + +// ds4_kernels.cu:268 hc_split_kernel - the sinkhorn normalisation of the hyper-connection mix. +// This is the SCALAR path of the CUDA kernel (lane 0 of the threadgroup does the token, a local 8x8 matrix), applied for every +// hc including 4: it computes the same mathematics as the warp-lane version, only without the shuffles. Correctness first - +// the lane-parallel hc == 4 version can be added behind a profile. +kernel void ds4_hc_split(device const float* mixes [[buffer(0)]], device const float* scale [[buffer(1)]], + device const float* base [[buffer(2)]], device float* pre [[buffer(3)]], + device float* post [[buffer(4)]], device float* comb [[buffer(5)]], + constant uint& hc [[buffer(6)]], constant uint& iters [[buffer(7)]], + constant float& eps [[buffer(8)]], uint b [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_threadgroup]]) { + if (lane != 0u) return; + const device float* mx = mixes + (size_t) b * 32u; + device float* pr = pre + (size_t) b * 8u; + device float* po = post + (size_t) b * 8u; + device float* cb = comb + (size_t) b * 64u; + for (uint i = 0; i < 8u; ++i) { + if (i >= hc) break; + pr[i] = 1.0f / (1.0f + exp(-(mx[i] * scale[0] + base[i]))) + eps; + po[i] = 2.0f / (1.0f + exp(-(mx[hc + i] * scale[1] + base[hc + i]))); + } + float c[8][8]; + for (uint j = 0; j < 8u; ++j) { + if (j >= hc) break; + float m = -3.402823466e+38f; + for (uint k = 0; k < 8u; ++k) { + if (k >= hc) break; + c[j][k] = mx[2u * hc + j * hc + k] * scale[2] + base[2u * hc + j * hc + k]; + m = fmax(m, c[j][k]); + } + float s = 0.0f; + for (uint k = 0; k < 8u; ++k) { + if (k >= hc) break; + c[j][k] = exp(c[j][k] - m); + s += c[j][k]; + } + for (uint k = 0; k < 8u; ++k) { + if (k >= hc) break; + c[j][k] = c[j][k] / s + eps; + } + } + for (uint k = 0; k < 8u; ++k) { + if (k >= hc) break; + float s = 0.0f; + for (uint j = 0; j < 8u; ++j) if (j < hc) s += c[j][k]; + for (uint j = 0; j < 8u; ++j) if (j < hc) c[j][k] /= s + eps; + } + for (uint it = 1; it < iters; ++it) { + for (uint j = 0; j < 8u; ++j) { + if (j >= hc) break; + float s = 0.0f; + for (uint k = 0; k < 8u; ++k) if (k < hc) s += c[j][k]; + for (uint k = 0; k < 8u; ++k) if (k < hc) c[j][k] /= s + eps; + } + for (uint k = 0; k < 8u; ++k) { + if (k >= hc) break; + float s = 0.0f; + for (uint j = 0; j < 8u; ++j) if (j < hc) s += c[j][k]; + for (uint j = 0; j < 8u; ++j) if (j < hc) c[j][k] /= s + eps; + } + } + for (uint j = 0; j < 8u; ++j) { + if (j >= hc) break; + for (uint k = 0; k < 8u; ++k) { if (k >= hc) break; cb[j * hc + k] = c[j][k]; } + } +} + +// ds4_kernels.cu:1938 moe_sort - count each virtual slot's picks, prefix-sum the counts, then fill in token order. +// kPicks = 8 (ds4_kernels.cuh:19): idx is [T][8], so q -> (t = q/8, i = q%8). +kernel void ds4_bounds_scan(device int* bounds [[buffer(0)]], constant uint& n_vs [[buffer(1)]], + uint tid [[thread_index_in_threadgroup]]) { + if (tid != 0u) return; + bounds[0] = 0; + for (uint s = 0; s < n_vs; ++s) bounds[s + 1] += bounds[s]; +} + +// simd_ballot returns a simd_vote, which does not cast to an integer in this toolchain; the same counts come from the +// prefix-sum and sum primitives, which give a lane its offset inside its simdgroup and the group its total. +static inline int ds4_warp_prefix(bool hit) { return simd_prefix_exclusive_sum(hit ? 1 : 0); } +static inline int ds4_warp_total(bool hit) { return simd_sum(hit ? 1 : 0); } + +kernel void ds4_moe_sort_count(device const int* idx [[buffer(0)]], constant uint& T [[buffer(1)]], + constant uint& kk [[buffer(2)]], device const int* vs [[buffer(3)]], + device int* bounds [[buffer(4)]], uint slot [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + threadgroup int warp_n[8]; + threadgroup int base_s; + const uint wp = tid / 32u, lane = tid % 32u; + if (tid == 0u) base_s = 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint n = T * kk; + for (uint q = tid; q < n; q += 256u) { + const uint t = q / kk, i = q - t * kk; + const bool hit = vs[idx[t * 8u + i]] == int(slot); + // one vote a warp: the running offset needs each warp's count, accumulated in slot order + const int wtot = ds4_warp_total(hit); + if ((tid % 32u) == 0u) warp_n[wp] = wtot; + threadgroup_barrier(mem_flags::mem_threadgroup); + int total = 0; + for (uint w = 0; w < 8u; ++w) total += warp_n[w]; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0u) base_s += total; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0u) bounds[slot + 1] = base_s; +} + +kernel void ds4_moe_sort_fill(device const int* idx [[buffer(0)]], constant uint& T [[buffer(1)]], + constant uint& kk [[buffer(2)]], device const int* vs [[buffer(3)]], + device int* bounds [[buffer(4)]], device int* row_tok [[buffer(5)]], + device int* pos [[buffer(6)]], uint slot [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + threadgroup int warp_n[8]; + threadgroup int base_s; + const uint wp = tid / 32u, lane = tid % 32u; + if (tid == 0u) base_s = bounds[slot]; + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint n = T * kk; + for (uint q = tid; q < n; q += 256u) { + const uint t = q / kk, i = q - t * kk; + const bool hit = vs[idx[t * 8u + i]] == int(slot); + const int wpre = ds4_warp_prefix(hit), wtot = ds4_warp_total(hit); + if ((tid % 32u) == 0u) warp_n[wp] = wtot; + threadgroup_barrier(mem_flags::mem_threadgroup); + int before = 0, total = 0; + for (uint w = 0; w < 8u; ++w) { if (w < wp) before += warp_n[w]; total += warp_n[w]; } + if (hit) { + const int r = base_s + before + wpre; + row_tok[r] = int(t); + pos[t * 8u + i] = r; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0u) base_s += total; + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +// ds4_kernels.cu:1690 engram_dequant_tok_kernel - one hash-table row a threadgroup: IQ4_NL blocks to F32. +// fmt 2 (this model's Engram): an 18-byte block is fp16 d then 16 packed nibbles; the low nibble of q[j] is value j and the +// high nibble value j + 16, each an index into ggml-common.h's kvalues_iq4nl (ds4_kernels.cu:1280). +constant int ds4_kIq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; + +kernel void ds4_engram_dequant(device const uchar* src [[buffer(0)]], constant ulong& per [[buffer(1)]], + constant uint& cols [[buffer(2)]], constant uint& hd [[buffer(3)]], + device float* out [[buffer(4)]], constant int& fmt [[buffer(5)]], + uint r [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const uint t = r / cols, c = r - t * cols; + const device uchar* rows = src + (size_t) t * per + (size_t) c * hd; + if (fmt == 2) { + for (uint e = tid; e < hd; e += 256u) { + const device uchar* b = rows + (size_t) (e / 32u) * 18u; + const uint q = b[2u + (e & 15u)]; + const int v = (e & 16u) ? int(q >> 4) : int(q & 15u); + const ushort dh = ushort(b[0]) | (ushort(b[1]) << 8); + out[(size_t) r * hd + e] = float(as_type(dh)) * float(ds4_kIq4nl[v]); + } + return; + } + // fmt 1 (Q6_K) and fmt 0 (E4M3 x scale) are not ported: write NaN, so a wrong path shows as NaN in the output rather than + // as a silent zero. + for (uint e = tid; e < hd; e += 256u) out[(size_t) r * hd + e] = NAN; +} + +// ds4_kernels.cu:341 hc_pre_norm_kernel - per token: a = sum_s mix[s]*h[token][s][e], one RMS over dim, then x = a*inv +// (times w when given). One threadgroup a token, 1024 threads; a thread holds at most 8 values (dim <= 8*1024 = 8192). +kernel void ds4_hc_pre_norm(device const float* h [[buffer(0)]], device const float* mix [[buffer(1)]], + device const float* w [[buffer(2)]], constant uint& dim [[buffer(3)]], + constant uint& hc [[buffer(4)]], constant float& eps [[buffer(5)]], + device float* x [[buffer(6)]], constant uint& has_w [[buffer(7)]], + uint b [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort wg [[simdgroup_index_in_threadgroup]]) { + threadgroup float red[32]; + const device float* hh = h + (size_t) b * hc * dim; + const device float* mm = mix + (size_t) b * 8u; + device float* xx = x + (size_t) b * dim; + float r[8]; + float ss = 0.0f; + uint n = 0; + for (uint e = tid; e < dim; e += 1024u, ++n) { + float a = 0.0f; + for (uint s = 0; s < 8u; ++s) if (s < hc) a += mm[s] * hh[(size_t) s * dim + e]; + if (n < 8u) r[n] = a; + ss += a * a; + } + // a 1024-thread reduction: 32 simdgroups, each into red[], then the first simdgroup sums those + float v = ds4_warp_sum(ss); + if (lane == 0u) red[wg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (wg == 0u) { + float t = lane < 32u ? red[lane] : 0.0f; + t = ds4_warp_sum(t); + if (lane == 0u) red[0] = t; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float inv = rsqrt(red[0] / float(dim) + eps); + n = 0; + for (uint e = tid; e < dim; e += 1024u, ++n) { + if (n >= 8u) break; + xx[e] = has_w != 0u ? r[n] * inv * w[e] : r[n]; + } +} + +// ds4_gguf.cu:186 q8_quant_kernel: a warp a 32-group - amax across the warp, sc = amax/127, q = rint(v/sc), d = sc. +// grid = (ceil((k/32) / 8 warps), rows), threadgroup = 256 (8 warps of 32). +kernel void ds4_q8_quant(device const float* x [[buffer(0)]], constant ulong& ldx [[buffer(1)]], + constant uint& k [[buffer(2)]], device char* q [[buffer(3)]], + device float* d [[buffer(4)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + const uint warp = tid / 32u, lane = tid % 32u; + const uint b = tgp.x * 8u + warp; // the 32-group index + const uint r = tgp.y; + if ((size_t) b * 32u >= (size_t) k) return; + const device float* xr = x + (size_t) r * ldx; + const float v = xr[(size_t) b * 32u + lane]; + const float am = simd_max(fabs(v)); + const float sc = am / 127.0f; + q[(size_t) r * k + (size_t) b * 32u + lane] = char(am > 0.0f ? rint(v / sc) : 0.0f); + if (lane == 0u) d[(size_t) r * (k / 32u) + b] = sc; +} + +// ds4_gguf.cu:371 mx_mv_kernel - the MXFP4 expert matrix-vector. +// grid = (ceil(n / 8 warps), jobs), threadgroup = 256. A warp a row; a lane strides the 32-value blocks; NC columns of +// activations share each decoded block (the T=8 window), then a warp sum and y[c][row] = alpha*v + bias[row]. +// Block layout (kMxBytes = 17 a 32 values): the E8M0 scale byte, then 16 bytes whose low nibble is value j and high nibble +// value j+16, each an index into the 16-entry codebook {0,1,2,3,4,6,8,12} and its negatives. +struct Ds4MxJob { + device const void* w; + ulong ld; // bytes a row + device const float* x; + ulong xcs; + device float* y; + ulong ycs; + device const float* bias; + float alpha; + int k, n, ncol; + device const char* xq; + device const float* xd; +}; + +static inline float ds4_e8m0(uint e) { // 2^(e - 128), ggml's GGML_E8M0_TO_FP32_HALF + return as_type(e < 2u ? (0x00200000u << e) : ((e - 1u) << 23)); +} + +static inline void ds4_mx_block(device const uchar* p, threadgroup const float* lut, thread float* w) { + const float d = ds4_e8m0(uint(p[0])); + for (uint j = 0; j < 16u; ++j) { + const uint q = uint(p[1u + j]); + w[j] = lut[q & 15u] * d; + w[j + 16u] = lut[q >> 4] * d; + } +} + +kernel void ds4_mx_mv(device const uchar* jobs [[buffer(0)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float lut[16]; + if (tid < 16u) { + const float v[8] = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 6.0f, 8.0f, 12.0f}; + lut[tid] = tid < 8u ? v[tid] : -v[tid - 8u]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const device Ds4MxJob& b = *(const device Ds4MxJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MxJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const device uchar* w = (const device uchar*) b.w + (size_t) row * b.ld; + thread float acc[8]; + for (uint c = 0; c < 8u; ++c) acc[c] = 0.0f; + const uint nblk = uint(b.k) / 32u; + for (uint kb = lane; kb < nblk; kb += 32u) { + float wv[32]; + ds4_mx_block(w + (size_t) kb * 17u, lut, wv); + for (uint c = 0; c < 8u; ++c) { + if (int(c) >= b.ncol) break; + const device float* x = b.x + c * b.xcs + (size_t) kb * 32u; + float a = acc[c]; + for (uint j = 0; j < 32u; ++j) a = fma(wv[j], x[j], a); + acc[c] = a; + } + } + for (uint c = 0; c < 8u; ++c) { + if (int(c) >= b.ncol) break; + float v = acc[c]; + v += simd_shuffle_xor(v, 16); + v += simd_shuffle_xor(v, 8); + v += simd_shuffle_xor(v, 4); + v += simd_shuffle_xor(v, 2); + v += simd_shuffle_xor(v, 1); + if (lane == 0u) b.y[(size_t) c * b.ycs + row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); + } +} + +// ds4_gguf.cu:320 q8_mv_kernel - the Q8_0 matvec against quantised activations. +// grid = (ceil(n / 8), jobs), threadgroup 256: a warp a row, a lane strides the 34-byte blocks (fp16 d + 32 int8), the dot is +// integer (CUDA's __dp4a in four-value groups; here the same sum written out), then acc += (dw * xd[c]) * sum. +// The staged-prefetch (ST) and vector variants of the CUDA kernel are prefetch strategies only - this reads the blocks directly. +kernel void ds4_q8_mv(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MxJob& b = *(const device Ds4MxJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MxJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const device uchar* w = (const device uchar*) b.w + (size_t) row * b.ld; + const uint nb = uint(b.k) / 32u; + float acc[8]; + for (uint c = 0; c < 8u; ++c) acc[c] = 0.0f; + for (uint kb = lane; kb < nb; kb += 32u) { + const device uchar* p = w + (size_t) kb * 34u; + const ushort dh = ushort(p[0]) | (ushort(p[1]) << 8); + const float dw = float(as_type(dh)); + int wv[32]; + for (uint j = 0; j < 32u; ++j) wv[j] = int((char) p[2u + j]); + for (uint c = 0; c < 8u; ++c) { + if (int(c) >= b.ncol) break; + const device char* xq = b.xq + (size_t) c * b.k + (size_t) kb * 32u; + int sum = 0; + for (uint j = 0; j < 32u; ++j) sum += wv[j] * int(xq[j]); + acc[c] = fma(dw * b.xd[(size_t) c * nb + kb], float(sum), acc[c]); + } + } + for (uint c = 0; c < 8u; ++c) { + if (int(c) >= b.ncol) break; + float v = acc[c]; + v += simd_shuffle_xor(v, 16); + v += simd_shuffle_xor(v, 8); + v += simd_shuffle_xor(v, 4); + v += simd_shuffle_xor(v, 2); + v += simd_shuffle_xor(v, 1); + if (lane == 0u) b.y[(size_t) c * b.ycs + row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); + } +} + +// ds4_kernels.cu:1092 ring_to_ext_kernel - one ring row of the KV cache to BF16: each value an E4M3 code times its group's +// E8M0 power of two (row[e] and row[d + e/32]); rows before the window start write zero. +static inline float ds4_e4m3(uint b) { // E4M3 (bias 7, no infinities, max 448): sign, 4-bit exp, 3-bit mantissa + const float s = (b & 0x80u) ? -1.0f : 1.0f; + const uint ex = (b >> 3) & 0xfu, m = b & 7u; + if (ex == 0u) return s * float(m) * (1.0f / 8.0f) * (1.0f / 64.0f); + if (ex == 15u && m == 7u) return NAN; + return s * exp2(float(ex) - 7.0f) * (1.0f + float(m) * 0.125f); +} +static inline float ds4_ue8m0(uint b) { return as_type(b << 23); } // 2^(b - 127) +static inline ushort ds4_f2bf16(float f) { + const uint u = as_type(f); + return ushort((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} + + +// ds4_kernels.cu:1104 ext_to_ring_kernel - the window's BF16 rows packed back into the E4M3 ring, each 32-value group at the +// smallest power-of-two scale it fits under (k), so v / 2^k is an E4M3 value times a power of two. +static inline uint ds4_e4m3_code(float v, int k) { + const uint sign = (v < 0.0f) ? 0x80u : 0u; + float a = fabs(v) * exp2(-float(k)); + if (a == 0.0f) return sign; + int e = int(floor(log2(a))); + int mant = int(round((a / exp2(float(e)) - 1.0f) * 8.0f)); + int ex = e + 7; + if (mant >= 8) { mant = 0; ++ex; } + if (ex <= 0) { mant = 0; ex = 0; } + if (ex > 15) { mant = 7; ex = 15; } + return sign | (uint(ex) << 3) | uint(mant & 7); +} + +kernel void ds4_ext_to_ring(device const ushort* ext [[buffer(0)]], device uchar* ring [[buffer(1)]], + device uchar* backup [[buffer(2)]], constant long& p0 [[buffer(3)]], + constant uint& n [[buffer(4)]], constant uint& T [[buffer(5)]], + constant uint& d [[buffer(6)]], constant ulong& rb [[buffer(7)]], constant uint& has_backup [[buffer(8)]], + uint t [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + if (T - t > n) return; // overwritten by a later row of the window anyway + device uchar* row = ring + (size_t) ((size_t(p0) + t) % n) * rb; + if (has_backup != 0u) { // (a bound buffer is never a null pointer here: the host says whether there is a backup) + for (size_t i = tid; i < rb; i += 256u) backup[(size_t) t * rb + i] = row[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const device ushort* src = ext + (size_t) (n + t) * d; + for (uint e0 = 0; e0 < d; e0 += 256u) { + const uint e = e0 + tid; + const float v = e < d ? as_type(uint(src[e]) << 16) : 0.0f; // (BF16 bits, not half) + uint a = as_type(v) & 0x7fffffffu; // magnitudes ordered as bits + for (uint o = 16u; o > 0u; o >>= 1u) a = max(a, simd_shuffle_xor(a, o)); + const int k = a == 0u ? 0 : int(a >> 23) - 135 + (((a & 0x7fffffu) > 0x600000u) ? 1 : 0); + if (e >= d) continue; + row[e] = uchar(ds4_e4m3_code(v, k)); + if ((e & 31u) == 0u) row[d + e / 32u] = uchar(k + 127); + } +} + + +// ds4_kernels.cu:916 unpack32 - 16 threads expand one MXFP4 KV row (288 B at d=512) into 32 BF16 values a thread: +// 16 code bytes at part*16, two E4M3 scale bytes at 256 + part*2, and the 4-bit code is itself an FP4 bit pattern +// (e2m1_val14: sign and exponent/mantissa bits spliced into a half). +static inline float ds4_e2m1_val14(uint c) { + const ushort h = (ushort) (((c & 8u) << 12) | ((c & 7u) << 9)); + return float(as_type(h)); +} + +// ds4_kernels.cu:937 unpack_sel_kernel - the selected rows of the compressed cache into the attention scratch, one row a +// group of 16 threads. n_sel comes either from the generated schedule (gen: min(kcap, (pos+1)/R)) or from nsel[t]. +kernel void ds4_unpack_sel(device const uchar* ckv [[buffer(0)]], device const int* sel [[buffer(1)]], + constant uint& sel_stride [[buffer(2)]], constant int& gen [[buffer(3)]], + constant long& p0 [[buffer(4)]], constant int& t0 [[buffer(5)]], + constant int& R [[buffer(6)]], constant int& kcap [[buffer(7)]], + device const int* nsel [[buffer(8)]], constant uint& cap [[buffer(9)]], + device ushort* out [[buffer(10)]], constant ulong& rb [[buffer(11)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const uint t = tgp.y; + int n_sel; + if (gen != 0) { + const long pa = p0 + long(t0) + long(t); + n_sel = R > 0 ? int(min(long(kcap), (pa + 1) / long(R))) : 0; + } else { + n_sel = nsel[t]; + } + const uint i = tgp.x * 16u + (tid >> 4u), part = tid & 15u; + if (int(i) >= n_sel) return; + const size_t r = (size_t) t * cap + i; + const device uchar* rr = ckv + (size_t) sel[(size_t) t * sel_stride + i] * rb; + const uint4 cw = *(const device uint4*) (rr + part * 16u); + const float k0 = ds4_e4m3(uint(rr[256u + part * 2u])) * 16384.0f; + const float k1 = ds4_e4m3(uint(rr[256u + part * 2u + 1u])) * 16384.0f; + const uint wv[4] = {cw.x, cw.y, cw.z, cw.w}; + uint o[16]; + for (uint k = 0; k < 16u; ++k) { + const uint b = wv[k >> 2] >> (8u * (k & 3u)); + const float ks = k < 8u ? k0 : k1; + const float lo = ds4_e2m1_val14(b) * ks, hi = ds4_e2m1_val14(b >> 4) * ks; + o[k] = (as_type(lo) >> 16) | (as_type(hi) & 0xffff0000u); + } + device uint* dst = (device uint*) (out + r * 512u + part * 32u); + for (uint q = 0; q < 4u; ++q) { + dst[q * 4u + 0u] = o[4u * q + 0u]; + dst[q * 4u + 1u] = o[4u * q + 1u]; + dst[q * 4u + 2u] = o[4u * q + 2u]; + dst[q * 4u + 3u] = o[4u * q + 3u]; + } +} + + +// ds4_kernels.cu:1772 moe_combine_kernel - a token's expert outputs back into its row: for each pick whose routed row is in +// [r0, r1), acc += w[t][i] * y[r - r0][e]; untouched when no pick falls in range. +kernel void ds4_moe_combine(device float* acc [[buffer(0)]], device const float* y [[buffer(1)]], + device const int* pos [[buffer(2)]], device const float* w [[buffer(3)]], + constant int& kk [[buffer(4)]], constant int& r0 [[buffer(5)]], + constant int& r1 [[buffer(6)]], constant uint& dim [[buffer(7)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const uint e = tgp.x * 256u + tid; + if (e >= dim) return; + const uint t = tgp.y; + float a = acc[(size_t) t * dim + e]; + bool any = false; + for (int i = 0; i < 8; ++i) { + if (i >= kk) break; + const int r = pos[(size_t) t * 8u + i]; + if (r < r0 || r >= r1) continue; + a = fma(w[(size_t) t * 8u + i], y[(size_t) (r - r0) * dim + e], a); + any = true; + } + if (any) acc[(size_t) t * dim + e] = a; +} + + +// ds4_kernels.cu:362 hc_post_kernel - like hc_post_inplace, but reads a separate residual and writes a separate output: +// out[t][k][e] = post[t][k]*y[t][e] + sum_j comb[t][j][k]*res[t][j][e] +kernel void ds4_hc_post_v(device const float* y [[buffer(0)]], device const float* res [[buffer(1)]], + device const float* post [[buffer(2)]], device const float* comb [[buffer(3)]], + device float* out [[buffer(4)]], constant uint& dim [[buffer(5)]], + constant uint& hc [[buffer(6)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + const uint i = tgp.x * 256u + tid; + if (i >= dim * hc) return; + const uint t = tgp.y, k = i / dim, e = i - k * dim; + const device float* r = res + (size_t) t * hc * dim; + float a = post[(size_t) t * 8u + k] * y[(size_t) t * dim + e]; + for (uint j = 0; j < 8u; ++j) { + if (j >= hc) break; + a += comb[(size_t) t * 64u + j * hc + k] * r[(size_t) j * dim + e]; + } + out[(size_t) t * hc * dim + i] = a; +} + + +// ds4_kernels.cu:441 gemv_f32_kernel - a warp a row of F32 weights: float4 loads of W[row] and x[t], four FMAs a step, +// lane-strided over k/4, then the warp sum into y[t][row]. +kernel void ds4_gemv_f32(device const float* W [[buffer(0)]], device const float* x [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& kk [[buffer(3)]], + device float* y [[buffer(4)]], constant uint& ldy [[buffer(5)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + const uint row = tgp.x * 8u + (tid >> 5); + if (row >= n) return; + const uint t = tgp.y; + const device float4* w = (const device float4*) (W + (size_t) row * kk); + const device float4* xv = (const device float4*) (x + (size_t) t * kk); + float acc = 0.0f; + for (uint c = lane; c < kk / 4u; c += 32u) { + const float4 a = w[c], b = xv[c]; + acc = fma(a.x, b.x, acc); + acc = fma(a.y, b.y, acc); + acc = fma(a.z, b.z, acc); + acc = fma(a.w, b.w, acc); + } + acc += simd_shuffle_xor(acc, 16); + acc += simd_shuffle_xor(acc, 8); + acc += simd_shuffle_xor(acc, 4); + acc += simd_shuffle_xor(acc, 2); + acc += simd_shuffle_xor(acc, 1); + if (lane == 0u) y[(size_t) t * ldy + row] = acc; +} + + +// ds4_kernels.cu:1353 route_publish_kernel - the MoE router: sc = sqrt(softplus(logit)), bs = sc + bias, then k rounds of a +// block-wide argmax (a strided max, a warp reduction, a 16-warp pass), the picks' sc summed into the weights, and each pick's +// device address published to plan_ptr. The prefetch requests (req) are NOT published here yet: that is the schedule, not the +// answer, and it can be turned on separately. +kernel void ds4_route_publish(device const float* logits [[buffer(0)]], device const float* bias [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& kk [[buffer(3)]], + constant float& route_scale [[buffer(4)]], device const ulong* vram [[buffer(5)]], + device int* idx [[buffer(6)]], device float* w [[buffer(7)]], + device ulong* plan_ptr [[buffer(8)]], device const float* bias_vl [[buffer(9)]], + constant uint& vl [[buffer(10)]], constant uint& has_vl [[buffer(11)]], + uint t [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + threadgroup float sc[1024]; + threadgroup float bs[1024]; + threadgroup float bv[16]; + threadgroup int bi[16]; + threadgroup int pick[8]; + const device float* lg = logits + (size_t) t * 1024u; + const device float* b = (has_vl != 0u && ((vl >> t) & 1u) != 0u) ? bias_vl : bias; + for (uint i = tid; i < n; i += 512u) { + const float l = lg[i]; + const float sp = l > 20.0f ? l : log(1.0f + exp(l)); + sc[i] = sqrt(sp); + bs[i] = sc[i] + b[i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint lane = tid & 31u, wp = tid >> 5u; + for (uint r = 0; r < 8u; ++r) { + if (r >= kk) break; + float best = -3.402823466e+38f; + int bix = 0x7fffffff; + for (uint i = tid; i < n; i += 512u) { + if (bs[i] > best || (bs[i] == best && int(i) < bix)) { best = bs[i]; bix = int(i); } + } + for (uint o = 16u; o > 0u; o >>= 1u) { + const float ov = simd_shuffle_xor(best, (ushort) o); + const int oi = simd_shuffle_xor(bix, (ushort) o); + if (ov > best || (ov == best && oi < bix)) { best = ov; bix = oi; } + } + if (lane == 0u && wp < 16u) { bv[wp] = best; bi[wp] = bix; } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0u) { + float bb = -3.402823466e+38f; + int ix = 0x7fffffff; + for (uint q = 0; q < 16u; ++q) { + if (bv[q] > bb || (bv[q] == bb && bi[q] < ix)) { bb = bv[q]; ix = bi[q]; } + } + pick[r] = ix; + bs[ix] = -3.402823466e+38f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0u) { + float sum = 0.0f; + for (uint r = 0; r < 8u; ++r) { if (r >= kk) break; sum += sc[pick[r]]; } + for (uint r = 0; r < 8u; ++r) { + if (r >= kk) break; + const uint p_ = uint(pick[r]); + const float wv = sc[p_] / (sum + 1e-20f) * route_scale; + idx[(size_t) t * 8u + r] = int(p_); + w[(size_t) t * 8u + r] = wv; + plan_ptr[(size_t) t * 8u + r] = vram[p_]; + } + } +} + + +// ds4_kernels.cu swiglu_kernel - the expert MLP's activation: h = silu(g) * u, with the clamp when lim > 0 +kernel void ds4_swiglu(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], + constant uint& n [[buffer(2)]], constant float& lim [[buffer(3)]], + device float* h [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + if (gid >= n) return; + float gv = g[gid], uv = u[gid]; + if (lim > 0.0f) { // both inputs clamped, then the gate itself + gv = fmin(gv, lim); + uv = fmin(fmax(uv, -lim), lim); + } + h[gid] = (gv / (1.0f + exp(-gv))) * uv; // silu(g) * u +} + +kernel void ds4_ring_to_ext(device const uchar* ring [[buffer(0)]], device ushort* ext [[buffer(1)]], + constant long& p0 [[buffer(2)]], constant uint& n [[buffer(3)]], + constant uint& d [[buffer(4)]], constant ulong& rb [[buffer(5)]], + uint r [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const long p = p0 - long(n) + long(r); + const device uchar* row = ring + (size_t) (p >= 0 ? (p % long(n)) : 0) * rb; + for (uint e = tid; e < d; e += 256u) { + ext[(size_t) r * d + e] = p >= 0 ? ds4_f2bf16(ds4_e4m3(uint(row[e])) * ds4_ue8m0(uint(row[d + e / 32u]))) : ushort(0); + } +} + +// ---------------------------------------------------------------- ds4_gguf.cu:577 mx_combine_kernel +// +// A token's expert outputs summed in plan order (an expert off the device adds 0), the CPU lane's part added when its +// flag is set, then the shared expert's - ds4_gguf.cu:169's combine(). Grid (n_embd / 128, T), 128 threads. +struct Ds4Comb { + device const ulong* plan_ptr; // [T][8]: nonzero when the device holds that expert + device const float* plan_w; // [T][8] + device const float* part; // [entries][n_embd] + device const float* cpu_part; // [T][cpu_stride], or null + device const int* cpu_flag; // [T], or null + device const float* sh_out; // [T][n_embd], or null + device float* out; + long cpu_stride, out_stride; + int k, n_embd; +}; + +kernel void ds4_mx_combine(constant Ds4Comb& a [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + const uint r = tgp.x * 128u + tid; + if (int(r) >= a.n_embd) return; + const int t = int(tgp.y); + float m = 0.0f; + for (int i = 0; i < a.k; ++i) { // kPicks <= 8 + const int e = t * 8 + i; + const float pi = a.plan_ptr[e] != 0ul ? a.part[(size_t) e * (size_t) a.n_embd + r] : 0.0f; + m = fma(a.plan_w[t * 8 + i], pi, m); + } + if (a.cpu_flag != nullptr && a.cpu_flag[t] != 0) m += a.cpu_part[(size_t) t * (size_t) a.cpu_stride + r]; + m += a.sh_out != nullptr ? a.sh_out[(size_t) t * (size_t) a.n_embd + r] : 0.0f; + a.out[(size_t) t * (size_t) a.out_stride + r] = m; +} + +// ---------------------------------------------------------------- ds4_kernels.cu:1286 engram_dequant_kernel / :1311 engram_gate_kernel +// +// The Engram tables' row decoder and gate (this model's tables are fmt 2: IQ4_NL, an 18-byte block of 32 values whose +// nibbles index ds4_kIq4nl above). +kernel void ds4_engram_dequant_rows(device const uchar* rows [[buffer(0)]], constant uint& hd [[buffer(1)]], + device float* out [[buffer(2)]], constant int& fmt [[buffer(3)]], + uint r [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + if (fmt == 2) { + const device uchar* row = rows + (size_t) r * (size_t) hd; + for (uint e = tid; e < hd; e += 256u) { + const device uchar* b = row + (size_t) (e / 32u) * 18u; + const uint q = b[2u + (e & 15u)]; + const int v = (e & 16u) ? int(q >> 4) : int(q & 15u); + const ushort dh = ushort(b[0]) | (ushort(b[1]) << 8); + out[(size_t) r * (size_t) hd + e] = float(as_type(dh)) * float(ds4_kIq4nl[v]); + } + return; + } + // fmt 1 (Q6_K) and fmt 0 (E4M3 x scale) are not ported: NaN rather than a silent zero (as ds4_engram_dequant does) + for (uint e = tid; e < hd; e += 256u) out[(size_t) r * (size_t) hd + e] = NAN; +} + +// ds4_kernels.cu's block_sum for a 1024-thread threadgroup: a warp sum a warp, then every warp sums the 32 partials (the +// same result in all of them, so no warp diverges at a barrier) +static inline float ds4_block_sum32(float v, threadgroup float* red, ushort lane, ushort w) { + v = ds4_warp_sum(v); + if (lane == 0) red[w] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = lane < 32u ? red[lane] : 0.0f; + t = ds4_warp_sum(t); + threadgroup_barrier(mem_flags::mem_threadgroup); + return t; +} + +kernel void ds4_engram_gate(device float* h [[buffer(0)]], device const float* kv [[buffer(1)]], + device const float* w [[buffer(2)]], constant uint& hc [[buffer(3)]], + constant uint& dim [[buffer(4)]], constant float& eps [[buffer(5)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort wg [[simdgroup_index_in_threadgroup]]) { + threadgroup float red[32]; + const uint k = tgp.x, t = tgp.y; + device float* hk = h + ((size_t) t * hc + (size_t) k) * (size_t) dim; + const device float* kvt = kv + (size_t) t * (size_t) (hc + 1u) * (size_t) dim; + const device float* key = kvt + (size_t) k * (size_t) dim; + const device float* val = kvt + (size_t) hc * (size_t) dim; + const device float* wk = w + (size_t) k * (size_t) dim; + float sh = 0.0f, sk = 0.0f, sd = 0.0f; + for (uint e = tid; e < dim; e += 1024u) { + const float hv = hk[e], kve = key[e]; + sh += hv * hv; + sk += kve * kve; + sd += hv * wk[e] * kve; + } + sh = ds4_block_sum32(sh, red, lane, wg); + sk = ds4_block_sum32(sk, red, lane, wg); + sd = ds4_block_sum32(sd, red, lane, wg); + const float rstd = rsqrt(sh / float(dim) + eps) * rsqrt(sk / float(dim) + eps); + const float dot = sd * rstd * rsqrt(float(dim)); + const float sq = copysign(sqrt(fmax(fabs(dot), 1e-6f)), dot); + const float gate = 1.0f / (1.0f + exp(-sq)); + for (uint e = tid; e < dim; e += 1024u) hk[e] = hk[e] + gate * val[e]; +} + +// ds4_gguf.cu:72 bf16_mv_kernel - a BF16 matrix's matvec against F32 activations (the engram WKV here, and the +// BF16 dense matrices the jobs' f == 0 names). Same frame as ds4_mx_mv above: a warp a row, a lane strides the row; here a +// lane 8 BF16 weights (16 bytes = one uint4) and 8 activations, so the row is walked in 256-value steps. +// A pair's F32 is bit-exact - the low half zeroed - which is bf16_mv_kernel's __uint_as_float(w << 16). +kernel void ds4_bf16_mv(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MxJob& b = *(const device Ds4MxJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MxJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const device ushort* w = (const device ushort*) b.w + (size_t) row * b.ld; + thread float acc[16]; // (the CUDA template's largest NC; ncol cuts the rest) + for (uint c = 0; c < 16u; ++c) acc[c] = 0.0f; + for (uint i = uint(lane) * 8u; i < uint(b.k); i += 256u) { + const uint4 q = *(const device uint4*) (w + i); + const float wf[8] = {as_type(q.x << 16), as_type(q.x & 0xffff0000u), + as_type(q.y << 16), as_type(q.y & 0xffff0000u), + as_type(q.z << 16), as_type(q.z & 0xffff0000u), + as_type(q.w << 16), as_type(q.w & 0xffff0000u)}; + for (uint c = 0; c < 16u; ++c) { + if (int(c) >= b.ncol) break; + const device float4* xp = (const device float4*) (b.x + c * b.xcs + (size_t) i); + const float4 x0 = xp[0], x1 = xp[1]; + float a = acc[c]; + a = fma(wf[0], x0.x, a); + a = fma(wf[1], x0.y, a); + a = fma(wf[2], x0.z, a); + a = fma(wf[3], x0.w, a); + a = fma(wf[4], x1.x, a); + a = fma(wf[5], x1.y, a); + a = fma(wf[6], x1.z, a); + a = fma(wf[7], x1.w, a); + acc[c] = a; + } + } + for (uint c = 0; c < 16u; ++c) { + if (int(c) >= b.ncol) break; + float v = acc[c]; + v += simd_shuffle_xor(v, 16); + v += simd_shuffle_xor(v, 8); + v += simd_shuffle_xor(v, 4); + v += simd_shuffle_xor(v, 2); + v += simd_shuffle_xor(v, 1); + if (lane == 0u) b.y[(size_t) c * b.ycs + row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); + } +} + +// ---------------------------------------------------------------- ds4_kernels.cu:524 compress_kernel / :582 index_score_kernel +// +// The compressor (ratio > 1 layers): each token's kv / score row into its slot, and every ratio'th token pooled by a +// per-value softmax over the slots, scaled by a row RMS and the norm weights. Grid (1), 512 threads (d <= 1024 values). +static inline float ds4_block_sum_n(float v, threadgroup float* red, ushort lane, ushort w, uint nw) { + v = ds4_warp_sum(v); + if (lane == 0) red[w] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = lane < nw ? red[lane] : 0.0f; + t = ds4_warp_sum(t); // (every warp sums the same partials: no divergence at a barrier) + threadgroup_barrier(mem_flags::mem_threadgroup); + return t; +} + +kernel void ds4_compress(device const float* kv [[buffer(0)]], device const float* score [[buffer(1)]], + device float* kv_state [[buffer(2)]], device float* score_state [[buffer(3)]], + constant int& ratio [[buffer(4)]], constant int& p0 [[buffer(5)]], constant int& T [[buffer(6)]], + constant uint& d [[buffer(7)]], device const float* wn [[buffer(8)]], + constant float& eps [[buffer(9)]], device float* latent [[buffer(10)]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], + ushort wg [[simdgroup_index_in_threadgroup]]) { + threadgroup float red[32]; + const float kfmax = -3.4028234663852886e+38f; // -FLT_MAX, as the CUDA writes it + int g = 0; + for (int t = 0; t < T; ++t) { + const int slot = (p0 + t) % ratio; + const bool pool = (p0 + t + 1) % ratio == 0; + for (uint e = tid; e < d; e += 512u) { + kv_state[(size_t) slot * d + e] = kv[(size_t) t * d + e]; + score_state[(size_t) slot * d + e] = score[(size_t) t * d + e]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (!pool) continue; + float val[2] = {0.0f, 0.0f}; // (d <= 2 * 512) + float ss = 0.0f; + uint n = 0; + for (uint e = tid; e < d; e += 512u, ++n) { + float m = kfmax; + for (int s = 0; s < ratio; ++s) m = fmax(m, score_state[(size_t) s * d + e]); + float den = 0.0f; + for (int s = 0; s < ratio; ++s) den += exp(score_state[(size_t) s * d + e] - m); + float acc = 0.0f; + for (int s = 0; s < ratio; ++s) + acc += kv_state[(size_t) s * d + e] * (exp(score_state[(size_t) s * d + e] - m) / den); + val[n] = acc; + ss += acc * acc; + } + const float inv = rsqrt(ds4_block_sum_n(ss, red, lane, wg, 16u) / float(d) + eps); + n = 0; + for (uint e = tid; e < d; e += 512u, ++n) latent[(size_t) g * d + e] = val[n] * inv * wn[e]; + ++g; + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +// ---------------------------------------------------------------- ds4_kernels.cu:582 index_score_kernel<128, KF> +// +// A thread a position: its key row (kfmt 3: 64 bytes of E2M1 codes (value 2i in the low nibble) then 4 UE8M0 scales; kfmt +// 0: BF16) against every head's query (a threadgroup tile), relu, then the heads' weighted sum. Grid (ceil(n / 256), T). +struct Ds4Rows { // ds4_kernels.cuh:115 Rows (kWin = 16) + int n[16]; + int k[16]; + int gen, R, kcap, ccap, blk; + long p0; +}; + +static inline void ds4_rows_at(constant Ds4Rows& r, int t, thread int& n, thread int& k) { + if (r.gen == 0) { + n = r.n[t]; + k = r.k[t]; + return; + } + const int nn = int((r.p0 + t + 1) / r.R); + const int cb = min(r.ccap, (nn + r.blk - 1) / r.blk); + switch (r.gen) { + case 1: n = nn; k = min(r.kcap, nn); break; + case 2: n = nn; k = cb; break; + case 3: n = cb * r.blk; k = min(r.kcap, nn); break; + default: n = (nn + r.blk - 1) / r.blk; k = r.kcap; break; + } +} + +kernel void ds4_index_score(device const float* q [[buffer(0)]], device const float* wq [[buffer(1)]], + device const uchar* kc [[buffer(2)]], device const int* cand [[buffer(3)]], + device float* out [[buffer(4)]], constant Ds4Rows& lim [[buffer(5)]], + constant int& heads [[buffer(6)]], constant int& n [[buffer(7)]], + constant int& cand_stride [[buffer(8)]], constant int& blk [[buffer(9)]], + constant long& out_stride [[buffer(10)]], constant int& kfmt [[buffer(11)]], + constant int& has_cand [[buffer(12)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + constexpr int HD = 128; + threadgroup float qs[32 * HD + 32]; + const int t = int(tgp.y); + for (uint i = tid; i < uint(heads) * uint(HD); i += 256u) qs[i] = q[(size_t) t * (size_t) heads * HD + i]; + for (uint i = tid; i < uint(heads); i += 256u) qs[uint(heads) * HD + i] = wq[(size_t) t * 64 + i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + const int i = int(tgp.x) * 256 + int(tid); + if (i >= n) return; + device float* o = out + (long) t * out_stride; + int ln, lk; + ds4_rows_at(lim, t, ln, lk); + int p = i; + if (has_cand != 0) { + if (i >= lk * blk) { + o[i] = -INFINITY; + return; + } + p = cand[(size_t) t * (size_t) cand_stride + (size_t) (i / blk)] * blk + i % blk; + } + if (p >= ln) { + o[i] = -INFINITY; + return; + } + const size_t kb = kfmt == 3 ? (size_t) ((HD / 2 + HD / 32 + 15) / 16 * 16) : (size_t) HD * 2; + const device uchar* kr = kc + (size_t) p * kb; + float acc[32]; + for (uint h = 0; h < 32u; ++h) acc[h] = 0.0f; + for (int c = 0; c < HD; c += 8) { + float kvv[8]; + if (kfmt == 3) { // 8 codes in 4 bytes, the group's power-of-two scale (2^(b - 127)) + const uint cw = uint(kr[c / 2]) | (uint(kr[c / 2 + 1]) << 8) | (uint(kr[c / 2 + 2]) << 16) | + (uint(kr[c / 2 + 3]) << 24); + const float ks = as_type(uint(kr[HD / 2 + c / 32]) << 23); + for (uint x = 0; x < 8u; ++x) kvv[x] = ds4_e2m1_val14(cw >> (4 * x)) * 16384.0f * ks; + } else { + const device ushort* b = (const device ushort*) (kr + (size_t) c * 2); + for (uint x = 0; x < 8u; ++x) kvv[x] = as_type(uint(b[x]) << 16); + } + for (uint h = 0; h < 32u; ++h) { + if (int(h) >= heads) break; + threadgroup const float* qh = qs + h * HD + c; + for (uint x = 0; x < 8u; ++x) acc[h] += qh[x] * kvv[x]; + } + } + float s = 0.0f; + for (uint h = 0; h < 32u; ++h) { + if (int(h) >= heads) break; + s += fmax(acc[h], 0.0f) * qs[uint(heads) * HD + h]; + } + o[i] = s; +} + +// The sparse-attention indexer with the GPU's matrix units. The scalar ds4_index_score spends 16.3 s of a 32K prompt +// (30 % of it) doing 32 head-dots a (query, key) pair one FMA at a time; this stages 16 fp4 keys and one query's q as +// halves and runs the 32 x 128 x 16 product through simdgroup_matrix, which measured 3.17-6.10x faster in the standalone +// harness (tools/ds4/idx_harness/idx_mma) with every output matching the scalar path to half precision (<= 7e-4 +// relative). Same arithmetic, same order per (query, key) up to that rounding. +// +// The fp4 weights are staged as e2m1 * ks (<= 12, half-safe) with the 16384 applied once in the epilogue - the full +// 16384 * ks reaches 2e5 and overflows half. +kernel void ds4_index_score_mma(device const float* q [[buffer(0)]], device const float* wq [[buffer(1)]], + device const uchar* kc [[buffer(2)]], device float* out [[buffer(4)]], + constant Ds4Rows& lim [[buffer(5)]], constant int& heads [[buffer(6)]], + constant int& n [[buffer(7)]], constant int& T [[buffer(8)]], + constant long& out_stride [[buffer(9)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + constexpr int HD = 128; + constexpr int KB = 80; // the fp4 row: 64 code bytes + 4 scale bytes, padded to 80 + constexpr int NK = 16; // keys a threadgroup scores (two key-blocks of 8) + threadgroup half qst[32 * HD]; // [head][dim] + threadgroup half kst[HD * NK]; // [dim][key], e2m1 * ks + threadgroup float part[4][NK]; // the four head-blocks' partial sums + threadgroup float accs[8][64]; // a simdgroup's 8x8 tile + const int k0 = int(tgp.x) * NK; + const int t = int(tgp.y); + if (t >= T) return; + for (uint i = tid; i < uint(heads) * uint(HD); i += 256u) qst[i] = half(q[(size_t) t * (size_t) heads * HD + i]); + for (int key = int(tid) / 8; key < NK; key += 32) { + const int idx = k0 + key; + if (idx >= n) continue; + const device uchar* kr = kc + (size_t) idx * (size_t) KB; + for (int c = 0; c < HD; c += 8) { + const uint cw = uint(kr[c / 2]) | (uint(kr[c / 2 + 1]) << 8) | (uint(kr[c / 2 + 2]) << 16) | + (uint(kr[c / 2 + 3]) << 24); + const float ks = as_type(uint(kr[HD / 2 + c / 32]) << 23); + for (uint x = 0; x < 8u; ++x) kst[(c + int(x)) * NK + key] = half(ds4_e2m1_val14(cw >> (4 * x)) * ks); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int hb = int(sg) / 2, kb = int(sg) % 2; // eight heads x eight keys a simdgroup + simdgroup_float8x8 acc = make_filled_simdgroup_matrix(0.0f); + for (int c = 0; c < HD; c += 8) { + simdgroup_half8x8 a, b; + simdgroup_load(a, qst, HD, ulong2(ulong(c), ulong(hb) * 8ul)); + simdgroup_load(b, kst, NK, ulong2(ulong(kb) * 8ul, ulong(c))); + simdgroup_multiply_accumulate(acc, a, b, acc); + } + simdgroup_store(acc, accs[sg], 8); + simdgroup_barrier(mem_flags::mem_threadgroup); + const uint lane = tid % 32u; + if (lane < 8u) { + const int key = kb * 8 + int(lane); + float ssum = 0.0f; + for (uint h = 0; h < 8u; ++h) { + const int hh = hb * 8 + int(h); + if (hh >= heads) break; + ssum += fmax(accs[sg][h * 8 + lane] * 16384.0f, 0.0f) * wq[(size_t) t * 64 + hh]; + } + part[hb][key] = ssum; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < uint(NK)) { + const int key = int(tid); + const int idx = k0 + key; + if (idx >= n) return; + int ln, lk; + ds4_rows_at(lim, t, ln, lk); + device float* o = out + (long) t * out_stride; + o[idx] = idx >= ln ? -INFINITY : (part[0][key] + part[1][key] + part[2][key] + part[3][key]); + } +} + +// ---------------------------------------------------------------- ds4_kernels.cu:142 finish_kernel (the KV pack) +// +// The same load, RMS and mrope as ds4_finish above, then the row written as the cache's packed codes (kv_row_bytes): the +// quantiser's scale byte a group (32 values, or 16 for the 4-bit MXFP4-style rows), the codes themselves from bit 0 up - +// for the 4-bit formats two values a byte, the second from the next lane. Written from the CUDA of the same name. + + + +static inline uint ds4_e2m1_code(float q) { // (bit 3 the sign, -0 kept; the magnitudes as 0 .. 7) + const float a = fabs(q); + const uint m = a <= 1.5f ? uint(a * 2.0f) : a <= 2.0f ? 4u : a <= 3.0f ? 5u : a <= 4.0f ? 6u : 7u; + return m | ((as_type(q) >> 31) << 3); +} + +kernel void ds4_finish_pack(device float* x [[buffer(0)]], constant long& stride [[buffer(1)]], + constant uint& len [[buffer(2)]], constant Ds4Finish& f [[buffer(3)]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], + ushort w [[simdgroup_index_in_threadgroup]], uint b [[threadgroup_position_in_grid]]) { + threadgroup float red[8]; + device float* v = x + (long) b * stride; + const uint nj = (len + 255u) / 256u; + float r[32]; + float ss = 0.0f; + for (uint j = 0; j < 32u; ++j) { + if (j >= nj) break; + const uint e = tid + j * 256u; + r[j] = e < len ? v[e] : 0.0f; + ss += r[j] * r[j]; + } + if (f.w != nullptr) { + const float tot = ds4_block_sum(ss, red, lane, w); + const float inv = rsqrt(tot / float(len) + f.eps); + for (uint j = 0; j < 32u; ++j) { + if (j >= nj) break; + const uint e = tid + j * 256u; + if (e < len) r[j] = r[j] * inv * f.w[e]; + } + } + if (f.rd > 0) { + const uint r0 = len - uint(f.rd); + const uint pos = uint(f.pos) + (b / uint(f.vpp)) * uint(f.pos_step); + for (uint j = 0; j < 32u; ++j) { + if (j >= nj) break; + const uint e = tid + j * 256u; + const float other = simd_shuffle_xor(r[j], 1); + if (e >= r0 && e < len) { + const uint i = (e - r0) >> 1; + const float ang = float(pos) * f.freqs[i]; + const float c = cos(ang); + const float sn = f.inverse ? -sin(ang) : sin(ang); + r[j] = ((e - r0) & 1u) == 0u ? r[j] * c - other * sn : other * sn + r[j] * c; + } + } + } + // the packed row: a group's scale and the group's codes as the format's (quant 1: d E4M3 codes at pr[e] and d/32 + // scale bytes at pr[len + e/32]; 2: E2M1 pairs at pr[e/2] with E4M3 scales at pr[len/2 + e/16]; 3: E2M1 pairs with + // UE8M0 scales at pr[len/2 + e/32]) + const uint g = f.quant == 2 ? 16u : 32u; + device uchar* pr = f.pdst + (long) b * f.pdst_stride; + for (uint j = 0; j < 32u; ++j) { + if (j >= nj) break; + const uint e = tid + j * 256u; + float a = e < len ? fabs(r[j]) : 0.0f; + for (uint o = g / 2u; o > 0u; o >>= 1) a = fmax(a, simd_shuffle_xor(a, ushort(o))); + float q = 0.0f, s = 1.0f; + if (e < len) { + if (f.quant == 1) { + s = ds4_pow2_ceil(fmax(a, 1e-4f) * (1.0f / 448.0f)); + q = ds4_e4m3_round(fmin(fmax(r[j] / s, -448.0f), 448.0f)); + } else if (f.quant == 2) { + s = ds4_e4m3_round(fmax(a, 6.0f * 0.001953125f) / 6.0f); + q = ds4_e2m1_rt(fmin(fmax(r[j] / s, -6.0f), 6.0f)); + } else { + s = ds4_pow2_ceil(fmax(a, 6.0f * 1.1754943508222875e-38f) * (1.0f / 6.0f)); + q = ds4_e2m1_rt(fmin(fmax(r[j] / s, -6.0f), 6.0f)); + } + r[j] = q * s; + } + if (f.quant == 1) { + if (e < len) pr[e] = uchar(ds4_e4m3_code(q, 0)); + } else { + const uint c = ds4_e2m1_code(q); + const uint hi = simd_shuffle_down(c, 1); // (value e + 1: the next lane's) + if (e < len && (e & 1u) == 0u) pr[e >> 1] = uchar(c | (hi << 4)); + } + if (e < len && e % g == 0u) + pr[(f.quant == 1 ? len : len / 2u) + e / g] = + uchar(f.quant == 2 ? ds4_e4m3_code(s, 0) : (as_type(s) >> 23)); + } + device ushort* d = f.dst != nullptr ? f.dst + (long) b * f.dst_stride : nullptr; + for (uint j = 0; j < 32u; ++j) { + if (j >= nj) break; + const uint e = tid + j * 256u; + if (e >= len) continue; + if (d != nullptr) { + const uint bits = as_type(r[j]); + d[e] = ushort((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16); // BF16 round to nearest even + } + } +} + +// ds4_kernels.cu:401 gemv_f16_rows_kernel - F16 weights against F32 activations (the indexer's wproj here): a warp a row +// of the matrix, a lane 8 halves (16 bytes, the uint4 above) and 8 activations a step, kRowsT = 2 tokens a block. Grid +// (ceil(n / 8), ceil(T / 2)), threadgroup 256. +kernel void ds4_gemv_f16_rows(device const ushort* W [[buffer(0)]], device const float* x [[buffer(1)]], + constant ulong& ldx [[buffer(2)]], constant int& n [[buffer(3)]], + constant int& k [[buffer(4)]], constant int& T [[buffer(5)]], + constant float& alpha [[buffer(6)]], device float* y [[buffer(7)]], + constant ulong& ldy [[buffer(8)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + constexpr int kRowsT = 2; + const int row = int(tgp.x) * 8 + int(tid >> 5); + const int t0 = int(tgp.y) * kRowsT, nt = min(kRowsT, T - t0); + if (row >= n) return; + const device ushort* w = W + (size_t) row * (size_t) k; + float acc[kRowsT]; + for (int t = 0; t < kRowsT; ++t) acc[t] = 0.0f; + for (int c = int(lane) * 8; c < k; c += 32 * 8) { + const uint4 q = *(const device uint4*) (w + c); + const float2 f0 = float2(as_type(q.x)), f1 = float2(as_type(q.y)); + const float2 f2 = float2(as_type(q.z)), f3 = float2(as_type(q.w)); + const float2 fs[4] = {f0, f1, f2, f3}; + float xv[kRowsT][8]; + for (int t = 0; t < kRowsT; ++t) + if (t < nt) { + const device float4* xp = (const device float4*) (x + (size_t) (t0 + t) * ldx + (size_t) c); + const float4 a4 = xp[0], b4 = xp[1]; + xv[t][0] = a4.x; + xv[t][1] = a4.y; + xv[t][2] = a4.z; + xv[t][3] = a4.w; + xv[t][4] = b4.x; + xv[t][5] = b4.y; + xv[t][6] = b4.z; + xv[t][7] = b4.w; + } + for (int i = 0; i < 4; ++i) { + const float2 f = fs[i]; + for (int t = 0; t < kRowsT; ++t) + if (t < nt) acc[t] += f.x * xv[t][2 * i] + f.y * xv[t][2 * i + 1]; + } + } + for (int t = 0; t < kRowsT; ++t) + if (t < nt) { + float a = acc[t]; + a += simd_shuffle_xor(a, 16); + a += simd_shuffle_xor(a, 8); + a += simd_shuffle_xor(a, 4); + a += simd_shuffle_xor(a, 2); + a += simd_shuffle_xor(a, 1); + if (lane == 0u) y[(size_t) (t0 + t) * ldy + row] = alpha * a; + } +} + +// ---------------------------------------------------------------- ds4_kernels.cu:650 topk_kernel +// +// A token's k largest scores: four 8-bit radix passes over the order keys of the row (a threadgroup of 1024, shared +// histograms), then the picks written in index order by a scan a pass, so the attention sums them in the same order every +// run and the ties kept are the lowest positions. +static inline uint ds4_ord_key(float f) { + const uint b = as_type(f); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +kernel void ds4_topk(device const float* vals [[buffer(0)]], constant long& stride [[buffer(1)]], + constant Ds4Rows& r [[buffer(2)]], device const int* cand [[buffer(3)]], + constant int& cand_stride [[buffer(4)]], constant int& blk [[buffer(5)]], + device int* out [[buffer(6)]], constant int& out_stride [[buffer(7)]], + constant int& has_cand [[buffer(8)]], uint tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], + ushort wid [[simdgroup_index_in_threadgroup]]) { + threadgroup uint hist[256]; + threadgroup uint prefix_s, mask_s; + threadgroup int remaining_s, cnt_gt, cnt_eq; + threadgroup int wgt[32], weq[32]; + const int t = int(tgp); + int n, k; + ds4_rows_at(r, t, n, k); + device const float* v = vals + (long) t * stride; + device int* o = out + (size_t) t * (size_t) out_stride; + device const int* cd = has_cand != 0 ? cand + (size_t) t * (size_t) cand_stride : nullptr; + if (n <= 0) return; + if (k >= n) { + for (uint i = tid; i < uint(n); i += 1024u) + o[i] = cd != nullptr ? cd[i / uint(blk)] * blk + int(i % uint(blk)) : int(i); + return; + } + if (tid == 0u) { + prefix_s = 0u; + mask_s = 0u; + remaining_s = k; + cnt_gt = 0; + cnt_eq = 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int shift = 24; shift >= 0; shift -= 8) { + for (uint i = tid; i < 256u; i += 1024u) hist[i] = 0u; + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint prefix = prefix_s, mask = mask_s; + for (int i = int(tid); i < n; i += 1024) { + const uint u = ds4_ord_key(v[i]); + if ((u & mask) == prefix) { + threadgroup atomic_uint* hp = (threadgroup atomic_uint*) &hist[(u >> shift) & 0xffu]; + atomic_fetch_add_explicit(hp, 1u, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0u) { + const int rem = remaining_s; + int cum = 0; + for (int d = 255; d >= 0; --d) { + const int h = int(hist[d]); + if (cum + h >= rem) { + prefix_s = prefix | (uint(d) << shift); + mask_s = mask | (0xffu << shift); + remaining_s = rem - cum; + break; + } + cum += h; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const uint thr = prefix_s; + const int remaining = remaining_s; + const int n_gt = k - remaining; + for (int i0 = 0; i0 < n; i0 += 1024) { + const int i = i0 + int(tid); + const uint u = i < n ? ds4_ord_key(v[i]) : 0u; + const bool gt = i < n && u > thr, eq = i < n && u == thr; + const auto bgv = simd_ballot(gt); + const auto bev = simd_ballot(eq); + const uint bg = uint((simd_vote::vote_t)bgv), be = uint((simd_vote::vote_t)bev); + if (lane == 0u) { + wgt[wid] = int(popcount(bg)); + weq[wid] = int(popcount(be)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (wid == 0u) { + const int g0 = wgt[lane], e0 = weq[lane]; + int gs = g0, es = e0; + for (uint o2 = 1u; o2 < 32u; o2 <<= 1) { + const int a = simd_shuffle_up(gs, ushort(o2)), b = simd_shuffle_up(es, ushort(o2)); + if (uint(lane) >= o2) { + gs += a; + es += b; + } + } + wgt[lane] = cnt_gt + gs - g0; + weq[lane] = cnt_eq + es - e0; + simdgroup_barrier(mem_flags::mem_none); + if (lane == 31u) { + cnt_gt += gs; + cnt_eq += es; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + int slot = -1; + const uint below = (1u << lane) - 1u; + if (gt) slot = wgt[wid] + int(popcount(bg & below)); + else if (eq) { + const int e = weq[wid] + int(popcount(be & below)); + if (e < remaining) slot = n_gt + e; + } + if (slot >= 0) o[slot] = cd != nullptr ? cd[i / blk] * blk + i % blk : i; + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +// ds4_kernels.cu:751 block_max_kernel - a block of `blk` positions' largest score (the indexer's candidate blocks), the +// newest position's block pinned in at +inf. Grid (ceil(nb_max / 256), T), 256 threads. +kernel void ds4_block_max(device const float* score [[buffer(0)]], constant long& stride [[buffer(1)]], + constant Ds4Rows& r [[buffer(2)]], constant int& blk [[buffer(3)]], + device float* bm [[buffer(4)]], constant long& bm_stride [[buffer(5)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const int b = int(tgp.x) * 256 + int(tid), t = int(tgp.y); + int n, k_; + ds4_rows_at(r, t, n, k_); + const int nb = (n + blk - 1) / blk; + if (b >= nb) return; + float m = -INFINITY; + for (int j = 0; j < blk; ++j) { + const int p = b * blk + j; + if (p < n) m = fmax(m, score[(long) t * stride + p]); + } + if (b == (n - 1) / blk) m = INFINITY; // the newest position's block is pinned in + bm[(long) t * bm_stride + b] = m; +} + +// ---------------------------------------------------------------- ds4_kernels.cu:1141 hc_mean_kernel / :1150 iota_rows_kernel +// +// The two small shapes the decode loop asks for: a token's mean over its hc streams, and a row's indices. +kernel void ds4_hc_mean(device const float* h [[buffer(0)]], constant int& hc [[buffer(1)]], + constant int& dim [[buffer(2)]], device float* out [[buffer(3)]], + constant long& out_stride [[buffer(4)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + const int e = int(tgp.x) * 256 + int(tid), t = int(tgp.y); + if (e >= dim) return; + float a = 0.0f; + for (int s = 0; s < hc; ++s) a += h[((size_t) t * (size_t) hc + (size_t) s) * (size_t) dim + (size_t) e]; + out[(long) t * out_stride + e] = a / float(hc); +} + +kernel void ds4_iota_rows(device int* out [[buffer(0)]], constant int& n [[buffer(1)]], constant int& stride [[buffer(2)]], + uint tg [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const uint t = tg; + for (uint i = tid; i < uint(n); i += 64u) out[(size_t) t * (size_t) stride + i] = int(i); +} + +// ---------------------------------------------------------------- vendor/sampler.cu:93 sampler_greedy_kernel +// +// The plain argmax over the vocabulary, ONE BLOCK (1024) a token: a thread's strided candidates with the repetition +// penalties applied (a threadgroup bitmap of the history window says which tokens can be penalised at all), a warp +// reduction, then warp 0's - a tie keeps the larger value and the SMALLER index, so this is the serial scan's answer. +kernel void ds4_sampler_greedy(device const float* logits [[buffer(0)]], constant int& n_vocab [[buffer(1)]], + device const int* history [[buffer(2)]], constant int& history_len [[buffer(3)]], + constant int& plen [[buffer(4)]], constant float& pen_repeat [[buffer(5)]], + constant float& pen_freq [[buffer(6)]], constant float& pen_present [[buffer(7)]], + device const float* logit_bias [[buffer(8)]], constant int& has_bias [[buffer(9)]], + device int* out [[buffer(10)]], uint tg [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], + ushort wid [[simdgroup_index_in_threadgroup]]) { + threadgroup uint penal_bits[4096]; // (n_vocab + 31) / 32 words (the host checks it fits) + threadgroup float sv[32]; + threadgroup int si[32]; + const int t = int(tg); + device const float* l = logits + (size_t) t * (size_t) n_vocab; + device const int* hrow = history != nullptr ? history + (size_t) t * (size_t) history_len : nullptr; + int hlen = 0; + if (hrow != nullptr) { + hlen = min(plen, history_len); + if (hlen < 0) hlen = 0; + hrow += history_len - hlen; // the window is the tail + } + const int bits_words = (n_vocab + 31) / 32; + const bool use_bits = hrow != nullptr && hlen > 0 && bits_words > 0; + if (use_bits) { + for (int w = int(tid); w < bits_words; w += 1024) penal_bits[w] = 0u; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int i = int(tid); i < hlen; i += 1024) + if (hrow[i] >= 0 && hrow[i] < n_vocab) { + threadgroup atomic_uint* pb = (threadgroup atomic_uint*) &penal_bits[hrow[i] >> 5]; + atomic_fetch_or_explicit(pb, 1u << (hrow[i] & 31), memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + float bv = -INFINITY; + int best = n_vocab; // (the "no candidate" index: it loses every comparison) + for (int v = int(tid); v < n_vocab; v += 1024) { + float s = has_bias != 0 ? l[v] + logit_bias[v] : l[v]; + int cnt = 0; + if (use_bits && (penal_bits[v >> 5] & (1u << (v & 31))) != 0u) + for (int i = 0; i < hlen; ++i) + if (hrow[i] == v) ++cnt; + if (cnt > 0) { + if (s <= 0.0f) s *= pen_repeat; + else s /= pen_repeat; + s -= float(cnt) * pen_freq + 1.0f * pen_present; + } + if (s > bv) { + bv = s; + best = v; + } + } + for (uint o = 16u; o > 0u; o >>= 1) { + const float ov = simd_shuffle_down(bv, ushort(o)); + const int oi = simd_shuffle_down(best, ushort(o)); + if (ov > bv || (ov == bv && oi < best)) { + bv = ov; + best = oi; + } + } + if (lane == 0u) { + sv[wid] = bv; + si[wid] = best; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (wid == 0u) { + float wv = sv[lane]; + int wi = si[lane]; + for (uint o = 16u; o > 0u; o >>= 1) { + const float ov = simd_shuffle_down(wv, ushort(o)); + const int oi = simd_shuffle_down(wi, ushort(o)); + if (ov > wv || (ov == wv && oi < wi)) { + wv = ov; + wi = oi; + } + } + if (lane == 0u) out[t] = (wi < n_vocab) ? wi : 0; // (a tie between two -inf candidates: 0) + } +} + +// ---------------------------------------------------------------- sparse attention (the prompt path) +// +// ds4_kernels.cu:954 sparse_attn_packed_kernel and :774 sparse_attn_kernel. The packed kernel stages sixteen +// rows of a tile in threadgroup memory (the window rows as they are, the compressed rows decoded to BF16) once for all +// G heads of a threadgroup, then does each head's dots, softmax and weighted sum. The BF16 kernel is the same arithmetic +// for row sets already stored as BF16 (a DSpark block); one head and query a threadgroup. +struct Ds4AttnRows { // ds4_kernels.cuh:149 AttnRows (kWin = 16) + int win0[16], nwin[16], nsel[16]; + int gen, t0, win, R, kcap; + long p0; +}; + +static inline void ds4_attn_counts(constant Ds4AttnRows& ar, int t, thread int& n_win, thread int& n_sel, thread int& w0) { + if (ar.gen != 0) { + const long ta = (long) ar.t0 + (long) t; + const long pa = ar.p0 + ta; + const long a = (long) ar.win - ar.p0, b = ta + 1; + const long first = a > b ? a : b; // max(win - p0, ta + 1) + w0 = int(first); + n_win = int((long) ar.win + ta - first + 1); + n_sel = ar.R > 0 ? int(((pa + 1) / (long) ar.R) < (long) ar.kcap ? (pa + 1) / (long) ar.R : (long) ar.kcap) : 0; + } else { + w0 = ar.win0[t]; + n_win = ar.nwin[t]; + n_sel = ar.nsel[t]; + } + if (n_win < 0) n_win = 0; + if (n_sel < 0) n_sel = 0; +} + +static inline float ds4_bf16_u32(uint v) { return as_type(v << 16); } + +// 32 values of a packed compressed row (kKvFp4, 512 values) as BF16 in four uint4s +static inline void ds4_attn_unpack32(const device uchar* rr, int part, threadgroup uint4* dst) { + const uint4 cw = *(const device uint4*) (rr + (size_t) part * 16u); + const float k0 = ds4_e4m3(uint(rr[256u + uint(part) * 2u])) * 16384.0f; + const float k1 = ds4_e4m3(uint(rr[256u + uint(part) * 2u + 1u])) * 16384.0f; + const uint wv[4] = {cw.x, cw.y, cw.z, cw.w}; + uint o[16]; + for (uint k = 0; k < 16u; ++k) { + const uint b = wv[k >> 2] >> (8u * (k & 3u)); + const float ks = k < 8u ? k0 : k1; + const float lo = ds4_e2m1_val14(b) * ks, hi = ds4_e2m1_val14(b >> 4) * ks; + o[k] = (as_type(lo) >> 16) | (as_type(hi) & 0xffff0000u); + } + dst[0] = uint4(o[0], o[1], o[2], o[3]); + dst[1] = uint4(o[4], o[5], o[6], o[7]); + dst[2] = uint4(o[8], o[9], o[10], o[11]); + dst[3] = uint4(o[12], o[13], o[14], o[15]); +} + +// rows [s0, s0 + n) of the tile at t0 into tile4; a thread stages one row's 32 values +static inline void ds4_attn_stage(device const ushort* wbase, device const uchar* ckv, device const int* sl, + int n_win, int t0, int s0, int n, int sr, int part, threadgroup uint4* tile4) { + if (sr >= n) return; + const int j = s0 + sr; + threadgroup uint4* dst = tile4 + (size_t) sr * 64u + (size_t) part * 4u; + const int jj = t0 + j; + if (jj < n_win) { + const device uint4* src = (const device uint4*) (wbase + (size_t) jj * 512u + (size_t) part * 32u); + dst[0] = src[0]; + dst[1] = src[1]; + dst[2] = src[2]; + dst[3] = src[3]; + return; + } + const device uchar* rr = ckv + (size_t) sl[jj - n_win] * 288u; + ds4_attn_unpack32(rr, part, dst); +} + +static inline float ds4_attn_warp_max(float v) { + v = fmax(v, simd_shuffle_xor(v, 16)); + v = fmax(v, simd_shuffle_xor(v, 8)); + v = fmax(v, simd_shuffle_xor(v, 4)); + v = fmax(v, simd_shuffle_xor(v, 2)); + v = fmax(v, simd_shuffle_xor(v, 1)); + return v; +} + +template +static inline void ds4_sparse_packed_run(device const float* q, device const ushort* win, device const uchar* ckv, + device const int* sel, uint sel_stride, constant Ds4AttnRows& ar, + device const float* sink, uint heads, float scale, device float* out, + threadgroup uint4* tile4, threadgroup float* sc, threadgroup float* m_run_s, + threadgroup float* tscale_s, threadgroup float* den_s, uint2 tgp, uint tid, + ushort lane_in, ushort wp_in) { + constexpr int W = 8 / G; // warps a head + const int hb = int(tgp.x) * G, t = int(tgp.y); + int n_win, n_sel, w0; + ds4_attn_counts(ar, t, n_win, n_sel, w0); + const device ushort* wbase = win + (size_t) w0 * 512u; + const device int* sl = sel + (size_t) t * (size_t) sel_stride; + const uint lane = uint(lane_in), wp = uint(wp_in); + const uint gw = wp / uint(W), wr = wp % uint(W); + float qr[16]; + for (int x = 0; x < 16; ++x) + qr[x] = q[((size_t) t * heads + uint(hb) + gw) * 512u + lane * 16u + uint(x)]; + if (tid < uint(G)) { + m_run_s[tid] = -1e30f; + den_s[tid] = 0.0f; + } + float o0[G], o1[G]; + for (int g = 0; g < G; ++g) o0[g] = o1[g] = 0.0f; + const int c = int(tid) * 2; + const int sr = int(tid) >> 4, part = int(tid) & 15; + threadgroup_barrier(mem_flags::mem_threadgroup); + const int total = n_win + n_sel; + for (int t0 = 0; t0 < total; t0 += 64) { + const int nt = min(64, total - t0); + for (int s0 = 0; s0 < nt; s0 += 16) { // the dots + const int n = min(16, nt - s0); + ds4_attn_stage(wbase, ckv, sl, n_win, t0, s0, n, sr, part, tile4); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int r = int(wr); r < n; r += W) { + const threadgroup uint4* rp = tile4 + (size_t) r * 64u + lane * 2u; + const uint4 a = rp[0], b = rp[1]; + const uint av[4] = {a.x, a.y, a.z, a.w}, bv[4] = {b.x, b.y, b.z, b.w}; + float s = 0.0f; + for (int x = 0; x < 4; ++x) { + s += qr[2 * x] * ds4_bf16_u32(av[x] & 0xffffu); + s += qr[2 * x + 1] * ds4_bf16_u32(av[x] >> 16); + } + for (int x = 0; x < 4; ++x) { + s += qr[8 + 2 * x] * ds4_bf16_u32(bv[x] & 0xffffu); + s += qr[8 + 2 * x + 1] * ds4_bf16_u32(bv[x] >> 16); + } + s = ds4_warp_sum(s); + if (lane == 0u) sc[uint(gw) * 64u + uint(s0 + r)] = s * scale; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (wp < uint(G)) { // head wp's softmax over the tile + threadgroup float* scg = sc + wp * 64u; + float m = -INFINITY; + for (int j = int(lane); j < nt; j += 32) m = fmax(m, scg[j]); + m = ds4_attn_warp_max(m); + const float m_old = m_run_s[wp]; + const float m_new = fmax(m_old, m); + float s = 0.0f; + for (int j = int(lane); j < nt; j += 32) { + const float p = exp(scg[j] - m_new); + scg[j] = p; + s += p; + } + s = ds4_warp_sum(s); + if (lane == 0u) { + tscale_s[wp] = exp(m_old - m_new); + den_s[wp] = den_s[wp] * tscale_s[wp] + s; + m_run_s[wp] = m_new; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int g = 0; g < G; ++g) { + const float cs = tscale_s[g]; + o0[g] *= cs; + o1[g] *= cs; + } + const threadgroup uint* tu = (const threadgroup uint*) tile4; + for (int s0 = 0; s0 < nt; s0 += 16) { // the weighted sums + const int n = min(16, nt - s0); + ds4_attn_stage(wbase, ckv, sl, n_win, t0, s0, n, sr, part, tile4); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int r = 0; r < n; ++r) { + const uint pr = tu[(size_t) r * 256u + uint(tid)]; + const float v0 = ds4_bf16_u32(pr & 0xffffu), v1 = ds4_bf16_u32(pr >> 16); + for (int g = 0; g < G; ++g) { + const float p = sc[uint(g) * 64u + uint(s0 + r)]; + o0[g] = fma(p, v0, o0[g]); + o1[g] = fma(p, v1, o1[g]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + for (int g = 0; g < G; ++g) { + const float den = den_s[g] + exp(sink[hb + g] - m_run_s[g]); + device float* ov = out + ((size_t) t * heads + uint(hb + g)) * 512u; + ov[c] = o0[g] / den; + ov[c + 1] = o1[g] / den; + } +} + +#define DS4_SPARSE_PACKED_SIG(NAME) \ + kernel void NAME(device const float* q [[buffer(0)]], device const ushort* win [[buffer(1)]], \ + device const uchar* ckv [[buffer(2)]], device const int* sel [[buffer(3)]], \ + constant uint& sel_stride [[buffer(4)]], constant Ds4AttnRows& ar [[buffer(5)]], \ + device const float* sink [[buffer(6)]], constant uint& heads [[buffer(7)]], \ + constant float& scale [[buffer(8)]], device float* out [[buffer(9)]], \ + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], \ + ushort lane [[thread_index_in_simdgroup]], ushort wp [[simdgroup_index_in_threadgroup]]) + +#define DS4_SPARSE_PACKED_BODY(G) \ + { \ + threadgroup uint4 tile4[16u * 64u]; \ + threadgroup float sc[8u * 64u]; \ + threadgroup float m_run_s[8u], tscale_s[8u], den_s[8u]; \ + ds4_sparse_packed_run(q, win, ckv, sel, sel_stride, ar, sink, heads, scale, out, tile4, sc, m_run_s, \ + tscale_s, den_s, tgp, tid, lane, wp); \ + } + +DS4_SPARSE_PACKED_SIG(ds4_sparse_attn_g8) DS4_SPARSE_PACKED_BODY(8) +DS4_SPARSE_PACKED_SIG(ds4_sparse_attn_g4) DS4_SPARSE_PACKED_BODY(4) +DS4_SPARSE_PACKED_SIG(ds4_sparse_attn_g2) DS4_SPARSE_PACKED_BODY(2) +DS4_SPARSE_PACKED_SIG(ds4_sparse_attn_g1) DS4_SPARSE_PACKED_BODY(1) + +#undef DS4_SPARSE_PACKED_SIG +#undef DS4_SPARSE_PACKED_BODY + +#ifndef DS4_MX_PLO +#define DS4_MX_PLO 0 // 1: the softmax weights as BF16 high + low (rms error vs the scalar kernel 2e-6 instead of 3e-4) +#endif +struct Ds4AttnRaw { uint4 v[4]; int kind; }; // a thread's 32 values of one row, as read: kind 0 window (v[0..3] BF16), 1 compressed (v[0] codes, v[1].x the two scale bytes), 2 none + +static inline void ds4_attn_ld(device const ushort* wbase, device const uchar* ckv, device const int* sl, int n_win, int t0, int n, + int sr, int pt, thread Ds4AttnRaw& r) { + r.kind = 2; + if (sr >= n) return; + const int jj = t0 + sr; + if (jj < n_win) { + const device uint4* src = (const device uint4*) (wbase + (size_t) jj * 512u + (size_t) pt * 32u); + r.v[0] = src[0]; r.v[1] = src[1]; r.v[2] = src[2]; r.v[3] = src[3]; + r.kind = 0; + } else { + const device uchar* rr = ckv + (size_t) sl[jj - n_win] * 288u; + r.v[0] = *(const device uint4*) (rr + (size_t) pt * 16u); + r.v[1] = uint4(uint(rr[256u + uint(pt) * 2u]), uint(rr[256u + uint(pt) * 2u + 1u]), 0u, 0u); + r.kind = 1; + } +} + +static inline void ds4_attn_st(thread const Ds4AttnRaw& r, int sr, int pt, threadgroup uint4* tile4) { + threadgroup uint4* dst = tile4 + (size_t) sr * 64u + (size_t) pt * 4u; + if (r.kind == 0) { + dst[0] = r.v[0]; dst[1] = r.v[1]; dst[2] = r.v[2]; dst[3] = r.v[3]; + } else if (r.kind == 1) { + const float k0 = ds4_e4m3(r.v[1].x) * 16384.0f; + const float k1 = ds4_e4m3(r.v[1].y) * 16384.0f; + const uint wv[4] = {r.v[0].x, r.v[0].y, r.v[0].z, r.v[0].w}; + uint o[16]; + for (uint k = 0; k < 16u; ++k) { + const uint b = wv[k >> 2] >> (8u * (k & 3u)); + const float ks = k < 8u ? k0 : k1; + const float lo = ds4_e2m1_val14(b) * ks, hi = ds4_e2m1_val14(b >> 4) * ks; + o[k] = (as_type(lo) >> 16) | (as_type(hi) & 0xffff0000u); + } + dst[0] = uint4(o[0], o[1], o[2], o[3]); + dst[1] = uint4(o[4], o[5], o[6], o[7]); + dst[2] = uint4(o[8], o[9], o[10], o[11]); + dst[3] = uint4(o[12], o[13], o[14], o[15]); + } else { + dst[0] = uint4(0); dst[1] = uint4(0); dst[2] = uint4(0); dst[3] = uint4(0); + } +} + +// ---------------------------------------------------------------- sparse attention on the matrix units (G = 8 heads a threadgroup) +// +// The same attention as ds4_sparse_attn_g8 (a token, eight heads, the window rows plus the selected compressed rows, softmax with the head's +// sink) but the dots and the weighted sum run on simdgroup 8x8 matrices: eight rows a pass are staged once as BF16 (exact), q and the +// softmax weights are split into a BF16 high and low part (so the products keep about 16 bits), accumulation is fp32, and the softmax is the +// online one (the running maximum rescales the accumulators through a diagonal matrix). +#ifndef DS4_MX_QK_LO +#define DS4_MX_QK_LO 0 // 0: the QK dot uses the BF16 high part of q only (faster, ~8 fewer bits) +#endif +// ATTN_PROBE (function constant 25): a timing probe. 1 skips the P.V accumulate, 2 skips the QK matmul - the numbers +// it produces are wrong on purpose; it exists to say which half of the pass the 20 s of a 32K prompt is in. +constant int ATTN_PROBE [[function_constant(25)]]; +constant int attn_probe = is_function_constant_defined(ATTN_PROBE) ? ATTN_PROBE : 0; + +kernel void ds4_sparse_attn_mx(device const float* q [[buffer(0)]], device const ushort* win [[buffer(1)]], + device const uchar* ckv [[buffer(2)]], device const int* sel [[buffer(3)]], + constant uint& sel_stride [[buffer(4)]], constant Ds4AttnRows& ar [[buffer(5)]], + device const float* sink [[buffer(6)]], constant uint& heads [[buffer(7)]], + constant float& scale [[buffer(8)]], device float* out [[buffer(9)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort sg [[simdgroup_index_in_threadgroup]]) { + threadgroup uint4 tile4[8u * 64u]; + threadgroup float qbuf[4096]; // q hi / lo as BF16 (8 KB each), later the output rows as float + threadgroup float part[8u * 64u]; + threadgroup float ssm[64]; + threadgroup bfloat phi[64], plo[64]; + threadgroup float dg[64]; + threadgroup float m_s[8], den_s[8], sc_s[8]; + threadgroup bfloat* qhi = (threadgroup bfloat*) qbuf; + threadgroup bfloat* qlo = qhi + 4096; + threadgroup const bfloat* kt = (threadgroup const bfloat*) tile4; + + const int hb = int(tgp.x) * 8, t = int(tgp.y); + int n_win, n_sel, w0; + ds4_attn_counts(ar, t, n_win, n_sel, w0); + const device ushort* wbase = win + (size_t) w0 * 512u; + const device int* sl = sel + (size_t) t * (size_t) sel_stride; + for (uint x = 0; x < 16u; ++x) { + const uint idx = tid * 16u + x, h = idx >> 9, d = idx & 511u; + const float v = q[((size_t) t * heads + uint(hb) + h) * 512u + d]; + const bfloat hi = bfloat(v); + qhi[idx] = hi; + if (DS4_MX_QK_LO) qlo[idx] = bfloat(v - float(hi)); // (unused when the QK dot takes the high part only) + } + if (tid < 8u) { + m_s[tid] = -1e30f; + den_s[tid] = 0.0f; + } + simdgroup_float8x8 O[8]; + for (int c = 0; c < 8; ++c) O[c] = make_filled_simdgroup_matrix(0.0f); + const int sr = int(tid) >> 4, pt = int(tid) & 15; + const int total = n_win + n_sel; + threadgroup_barrier(mem_flags::mem_threadgroup); + Ds4AttnRaw raw; + raw.kind = 2; + if (sr < 8) ds4_attn_ld(wbase, ckv, sl, n_win, 0, min(8, total), sr, pt, raw); + for (int t0 = 0; t0 < total; t0 += 8) { + const int n = min(8, total - t0); + if (sr < 8 && attn_probe != 5) ds4_attn_st(raw, sr, pt, tile4); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sr < 8) ds4_attn_ld(wbase, ckv, sl, n_win, t0 + 8, min(8, total - t0 - 8), sr, pt, raw); // the next pass's rows, in flight during this one + if (attn_probe != 2 && attn_probe != 6) { // S = q . K^T over this simdgroup's eighth of the 512 values + simdgroup_float8x8 acc = make_filled_simdgroup_matrix(0.0f); + for (int kb = int(sg) * 8; kb < int(sg) * 8 + 8; ++kb) { + simdgroup_bfloat8x8 a_hi, a_lo, b; + simdgroup_load(b, kt, 512, ulong2(ulong(kb) * 8ul, 0ul), true); + simdgroup_load(a_hi, qhi, 512, ulong2(ulong(kb) * 8ul, 0ul)); + simdgroup_load(a_lo, qlo, 512, ulong2(ulong(kb) * 8ul, 0ul)); + simdgroup_multiply_accumulate(acc, a_hi, b, acc); + if (DS4_MX_QK_LO) simdgroup_multiply_accumulate(acc, a_lo, b, acc); // (the low part of q: exactness vs speed) + } + simdgroup_store(acc, part + uint(sg) * 64u, 8); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 64u && attn_probe != 4 && attn_probe != 6) { + float s = 0.0f; + for (uint g = 0; g < 8u; ++g) s += part[g * 64u + tid]; + ssm[tid] = s * scale; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (attn_probe != 3) { // head sg's online softmax over the pass's rows + const bool valid = int(lane) < n; + const float s = valid ? ssm[uint(sg) * 8u + uint(lane)] : -INFINITY; + const float m = ds4_attn_warp_max(s); + const float m_old = m_s[sg]; + const float m_new = fmax(m_old, m); + const float p = valid ? exp(s - m_new) : 0.0f; + const float ps = ds4_warp_sum(p); + if (lane < 8u) { + const bfloat hi = bfloat(p); + phi[uint(sg) * 8u + uint(lane)] = hi; + plo[uint(sg) * 8u + uint(lane)] = bfloat(p - float(hi)); + } + if (lane == 0u) { + const float cs = exp(m_old - m_new); + sc_s[sg] = cs; + den_s[sg] = den_s[sg] * cs + ps; + m_s[sg] = m_new; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 64u) dg[tid] = (tid >> 3) == (tid & 7u) ? sc_s[tid >> 3] : 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (attn_probe != 1 && attn_probe != 6) { // O = diag(rescale) O + P V on this simdgroup's 64 columns + simdgroup_float8x8 D; + simdgroup_bfloat8x8 ph, pl; + simdgroup_load(D, dg, 8); + simdgroup_load(ph, phi, 8); + simdgroup_load(pl, plo, 8); + for (int c = 0; c < 8; ++c) { + simdgroup_float8x8 tmp; + simdgroup_multiply(tmp, D, O[c]); + simdgroup_bfloat8x8 v; + simdgroup_load(v, kt, 512, ulong2(ulong(sg) * 64ul + ulong(c) * 8ul, 0ul)); + simdgroup_multiply_accumulate(tmp, ph, v, tmp); +#if DS4_MX_PLO + simdgroup_multiply_accumulate(tmp, pl, v, tmp); +#endif + O[c] = tmp; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (int c = 0; c < 8; ++c) simdgroup_store(O[c], qbuf, 512, ulong2(ulong(sg) * 64ul + ulong(c) * 8ul, 0ul)); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint x = 0; x < 16u; ++x) { + const uint idx = tid * 16u + x, h = idx >> 9, d = idx & 511u; + const float den = den_s[h] + exp(sink[uint(hb) + h] - m_s[h]); + out[((size_t) t * heads + uint(hb) + h) * 512u + d] = qbuf[idx] / den; + } +} + +// TWO QUERIES A THREADGROUP, written out in full. The staged keys are shared; each query keeps its own accumulators, +// its own online-softmax state and its own key count (the causal shape makes the second query's set a superset, so the +// union staged is the second query's own set and the first masks its tail). +kernel void ds4_sparse_attn_mx2(device const float* q [[buffer(0)]], device const ushort* win [[buffer(1)]], + device const uchar* ckv [[buffer(2)]], device const int* sel [[buffer(3)]], + constant uint& sel_stride [[buffer(4)]], constant Ds4AttnRows& ar [[buffer(5)]], + device const float* sink [[buffer(6)]], constant uint& heads [[buffer(7)]], + constant float& scale [[buffer(8)]], device float* out [[buffer(9)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort sg [[simdgroup_index_in_threadgroup]]) { + threadgroup uint4 tile4[8u * 64u]; + threadgroup float qbuf[4096]; // two queries' q as BF16 (8 KB each), later a query's output rows + threadgroup float part[8u * 64u]; + threadgroup float ssm[64]; + threadgroup bfloat phi[64], plo[64]; + threadgroup float dg[64]; + threadgroup float m_s[16], den_s[16], sc_s[16]; + threadgroup bfloat* qhiA = (threadgroup bfloat*) qbuf; + threadgroup bfloat* qhiB = qhiA + 4096; // a query's q is 4096 bf16 + threadgroup float* obuf = qbuf; // 8 heads x 512 floats: qbuf again, once the key loop is done + threadgroup const bfloat* kt = (threadgroup const bfloat*) tile4; + + const int hb = int(tgp.x) * 8; + const int tA = int(tgp.y) * 2; + const int rows_total = (ar.gen == 0 && ar.t0 > 0) ? ar.t0 : 0x7fffffff; + const int nq = (tA + 1 < rows_total) ? 2 : 1; + int nwA, nsA, w0A, nwB, nsB, w0B; + ds4_attn_counts(ar, tA, nwA, nsA, w0A); + nwB = nwA; nsB = nsA; w0B = w0A; + if (nq > 1) ds4_attn_counts(ar, tA + 1, nwB, nsB, w0B); + const int totalA = nwA + nsA; + const int totalB = (nq > 1) ? nwB + nsB : totalA; + const int total = totalA > totalB ? totalA : totalB; + const device ushort* wbase = win + (size_t) w0A * 512u; + const device int* sl = sel + (size_t) tA * (size_t) sel_stride; + for (uint x = 0; x < 16u; ++x) { + const uint idx = tid * 16u + x, h = idx >> 9, d = idx & 511u; + qhiA[idx] = bfloat(q[((size_t) tA * heads + uint(hb) + h) * 512u + d]); + if (nq > 1) qhiB[idx] = bfloat(q[((size_t) (tA + 1) * heads + uint(hb) + h) * 512u + d]); + } + if (tid < 16u) { + m_s[tid] = -1e30f; + den_s[tid] = 0.0f; + } + simdgroup_float8x8 OA[8], OB[8]; + for (int c = 0; c < 8; ++c) { + OA[c] = make_filled_simdgroup_matrix(0.0f); + OB[c] = make_filled_simdgroup_matrix(0.0f); + } + const int sr = int(tid) >> 4, pt = int(tid) & 15; + threadgroup_barrier(mem_flags::mem_threadgroup); + Ds4AttnRaw raw; + raw.kind = 2; + if (sr < 8) ds4_attn_ld(wbase, ckv, sl, nwB, 0, min(8, total), sr, pt, raw); + for (int t0 = 0; t0 < total; t0 += 8) { + const int np = min(8, total - t0); + (void) np; + if (sr < 8) ds4_attn_st(raw, sr, pt, tile4); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sr < 8) ds4_attn_ld(wbase, ckv, sl, nwB, t0 + 8, min(8, total - t0 - 8), sr, pt, raw); + { // query A: S = q . K^T + simdgroup_float8x8 acc = make_filled_simdgroup_matrix(0.0f); + for (int kb = int(sg) * 8; kb < int(sg) * 8 + 8; ++kb) { + simdgroup_bfloat8x8 a_hi, b; + simdgroup_load(b, kt, 512, ulong2(ulong(kb) * 8ul, 0ul), true); + simdgroup_load(a_hi, qhiA, 512, ulong2(ulong(kb) * 8ul, 0ul)); + simdgroup_multiply_accumulate(acc, a_hi, b, acc); + } + simdgroup_store(acc, part + uint(sg) * 64u, 8); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 64u) { + float s = 0.0f; + for (uint g = 0; g < 8u; ++g) s += part[g * 64u + tid]; + ssm[tid] = s * scale; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + { // query A: online softmax + const int nq_here = min(8, totalA - t0); + const bool valid = int(lane) < nq_here; + const float s = valid ? ssm[uint(sg) * 8u + uint(lane)] : -INFINITY; + const float m = ds4_attn_warp_max(s); + const float m_old = m_s[0 + sg]; + const float m_new = fmax(m_old, m); + const float p = valid ? exp(s - m_new) : 0.0f; + const float ps = ds4_warp_sum(p); + if (lane < 8u) { + const bfloat hi = bfloat(p); + phi[uint(sg) * 8u + uint(lane)] = hi; + plo[uint(sg) * 8u + uint(lane)] = bfloat(p - float(hi)); + } + if (lane == 0u) { + const float cs = exp(m_old - m_new); + sc_s[0 + sg] = cs; + den_s[0 + sg] = den_s[0 + sg] * cs + ps; + m_s[0 + sg] = m_new; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 64u) dg[tid] = (tid >> 3) == (tid & 7u) ? sc_s[0 + (tid >> 3)] : 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + { // query A: O = diag(rescale) O + P V + simdgroup_float8x8 D; + simdgroup_bfloat8x8 ph, pl; + simdgroup_load(D, dg, 8); + simdgroup_load(ph, phi, 8); + simdgroup_load(pl, plo, 8); + for (int c = 0; c < 8; ++c) { + simdgroup_float8x8 tmp; + simdgroup_multiply(tmp, D, OA[c]); + simdgroup_bfloat8x8 v; + simdgroup_load(v, kt, 512, ulong2(ulong(sg) * 64ul + ulong(c) * 8ul, 0ul)); + simdgroup_multiply_accumulate(tmp, ph, v, tmp); +#if DS4_MX_PLO + simdgroup_multiply_accumulate(tmp, pl, v, tmp); +#endif + OA[c] = tmp; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (nq > 1) { + { // query B: S = q . K^T + simdgroup_float8x8 acc = make_filled_simdgroup_matrix(0.0f); + for (int kb = int(sg) * 8; kb < int(sg) * 8 + 8; ++kb) { + simdgroup_bfloat8x8 a_hi, b; + simdgroup_load(b, kt, 512, ulong2(ulong(kb) * 8ul, 0ul), true); + simdgroup_load(a_hi, qhiB, 512, ulong2(ulong(kb) * 8ul, 0ul)); + simdgroup_multiply_accumulate(acc, a_hi, b, acc); + } + simdgroup_store(acc, part + uint(sg) * 64u, 8); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 64u) { + float s = 0.0f; + for (uint g = 0; g < 8u; ++g) s += part[g * 64u + tid]; + ssm[tid] = s * scale; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + { // query B: online softmax + const int nq_here = min(8, totalB - t0); + const bool valid = int(lane) < nq_here; + const float s = valid ? ssm[uint(sg) * 8u + uint(lane)] : -INFINITY; + const float m = ds4_attn_warp_max(s); + const float m_old = m_s[8 + sg]; + const float m_new = fmax(m_old, m); + const float p = valid ? exp(s - m_new) : 0.0f; + const float ps = ds4_warp_sum(p); + if (lane < 8u) { + const bfloat hi = bfloat(p); + phi[uint(sg) * 8u + uint(lane)] = hi; + plo[uint(sg) * 8u + uint(lane)] = bfloat(p - float(hi)); + } + if (lane == 0u) { + const float cs = exp(m_old - m_new); + sc_s[8 + sg] = cs; + den_s[8 + sg] = den_s[8 + sg] * cs + ps; + m_s[8 + sg] = m_new; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 64u) dg[tid] = (tid >> 3) == (tid & 7u) ? sc_s[8 + (tid >> 3)] : 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + { // query B: O = diag(rescale) O + P V + simdgroup_float8x8 D; + simdgroup_bfloat8x8 ph, pl; + simdgroup_load(D, dg, 8); + simdgroup_load(ph, phi, 8); + simdgroup_load(pl, plo, 8); + for (int c = 0; c < 8; ++c) { + simdgroup_float8x8 tmp; + simdgroup_multiply(tmp, D, OB[c]); + simdgroup_bfloat8x8 v; + simdgroup_load(v, kt, 512, ulong2(ulong(sg) * 64ul + ulong(c) * 8ul, 0ul)); + simdgroup_multiply_accumulate(tmp, ph, v, tmp); +#if DS4_MX_PLO + simdgroup_multiply_accumulate(tmp, pl, v, tmp); +#endif + OB[c] = tmp; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + for (int c = 0; c < 8; ++c) simdgroup_store(OA[c], obuf, 512, ulong2(ulong(sg) * 64ul + ulong(c) * 8ul, 0ul)); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint x = 0; x < 16u; ++x) { + const uint idx = tid * 16u + x, h = idx >> 9, d = idx & 511u; + out[((size_t) tA * heads + uint(hb) + h) * 512u + d] = obuf[idx] / (den_s[h] + exp(sink[uint(hb) + h] - m_s[h])); + } + if (nq > 1) { + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int c = 0; c < 8; ++c) simdgroup_store(OB[c], obuf, 512, ulong2(ulong(sg) * 64ul + ulong(c) * 8ul, 0ul)); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint x = 0; x < 16u; ++x) { + const uint idx = tid * 16u + x, h = idx >> 9, d = idx & 511u; + out[((size_t) (tA + 1) * heads + uint(hb) + h) * 512u + d] = + obuf[idx] / (den_s[8u + h] + exp(sink[uint(hb) + h] - m_s[8u + h])); + } + } +} + +// kContig == false (the row set is BF16 already): one head and query a threadgroup, a warp a row's dot and warp 0 the +// softmax, every thread two output values. +static inline void ds4_sparse_bf16_run(device const float* q, device const ushort* win, device const ushort* ckv, + device const int* sel, uint sel_stride, constant Ds4AttnRows& ar, + device const float* sink, uint heads, float scale, device float* out, + threadgroup float* qs, threadgroup float* sc, threadgroup int* rowix, + threadgroup float* m_run_s, threadgroup float* tscale_s, + threadgroup float* den_s, uint2 tgp, uint tid, ushort lane_in, + ushort wp_in) { + const int h = int(tgp.x), t = int(tgp.y); + int n_win, n_sel, w0; + ds4_attn_counts(ar, t, n_win, n_sel, w0); + const device ushort* wbase = win + (size_t) w0 * 512u; + const device int* sl = sel + (size_t) t * (size_t) sel_stride; + const uint lane = uint(lane_in), wp = uint(wp_in); + for (uint i = tid; i < 512u; i += 256u) qs[i] = q[((size_t) t * heads + uint(h)) * 512u + i]; + if (tid == 0u) { + m_run_s[0] = -1e30f; + den_s[0] = 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + float o0 = 0.0f, o1 = 0.0f; + const int c = int(tid) * 2; + const int total = n_win + n_sel; + for (int t0 = 0; t0 < total; t0 += 64) { + const int nt = min(64, total - t0); + if (tid < uint(nt)) { + const int j = t0 + int(tid); + rowix[tid] = j < n_win ? -1 - j : sl[j - n_win]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int jj = int(wp); jj < nt; jj += 8) { + const int code = rowix[jj]; + const device ushort* row = code < 0 ? wbase + (size_t) (-1 - code) * 512u + : ckv + (size_t) code * 512u; + const device uint4* rp = (const device uint4*) (row + lane * 16u); + const uint4 a = rp[0], b = rp[1]; + const uint av[4] = {a.x, a.y, a.z, a.w}, bv[4] = {b.x, b.y, b.z, b.w}; + float s = 0.0f; + for (int x = 0; x < 4; ++x) { + s += qs[lane * 16u + uint(2 * x)] * ds4_bf16_u32(av[x] & 0xffffu); + s += qs[lane * 16u + uint(2 * x + 1)] * ds4_bf16_u32(av[x] >> 16); + } + for (int x = 0; x < 4; ++x) { + s += qs[lane * 16u + uint(8 + 2 * x)] * ds4_bf16_u32(bv[x] & 0xffffu); + s += qs[lane * 16u + uint(9 + 2 * x)] * ds4_bf16_u32(bv[x] >> 16); + } + s = ds4_warp_sum(s); + if (lane == 0u) sc[jj] = s * scale; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (wp == 0u) { + float m = -INFINITY; + for (int j = int(lane); j < nt; j += 32) m = fmax(m, sc[j]); + m = ds4_attn_warp_max(m); + const float m_old = m_run_s[0]; + const float m_new = fmax(m_old, m); + float s = 0.0f; + for (int j = int(lane); j < nt; j += 32) { + const float p = exp(sc[j] - m_new); + sc[j] = p; + s += p; + } + s = ds4_warp_sum(s); + if (lane == 0u) { + tscale_s[0] = exp(m_old - m_new); + den_s[0] = den_s[0] * tscale_s[0] + s; + m_run_s[0] = m_new; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float cs = tscale_s[0]; + o0 *= cs; + o1 *= cs; + for (int j = 0; j < nt; ++j) { + const int code = rowix[j]; + const device ushort* row = code < 0 ? wbase + (size_t) (-1 - code) * 512u + : ckv + (size_t) code * 512u; + const uint pr = *(const device uint*) (row + c); + o0 = fma(sc[j], ds4_bf16_u32(pr & 0xffffu), o0); + o1 = fma(sc[j], ds4_bf16_u32(pr >> 16), o1); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float den = den_s[0] + exp(sink[h] - m_run_s[0]); + device float* ov = out + ((size_t) t * heads + uint(h)) * 512u; + ov[c] = o0 / den; + ov[c + 1] = o1 / den; +} + +kernel void ds4_sparse_attn_bf16(device const float* q [[buffer(0)]], device const ushort* win [[buffer(1)]], + device const ushort* ckv [[buffer(2)]], device const int* sel [[buffer(3)]], + constant uint& sel_stride [[buffer(4)]], constant Ds4AttnRows& ar [[buffer(5)]], + device const float* sink [[buffer(6)]], constant uint& heads [[buffer(7)]], + constant float& scale [[buffer(8)]], device float* out [[buffer(9)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort wp [[simdgroup_index_in_threadgroup]]) { + threadgroup float qs[512]; + threadgroup float sc[64]; + threadgroup int rowix[64]; + threadgroup float m_run_s, tscale_s, den_s; + ds4_sparse_bf16_run(q, win, ckv, sel, sel_stride, ar, sink, heads, scale, out, qs, sc, rowix, &m_run_s, + &tscale_s, &den_s, tgp, tid, lane, wp); +} + +// ds4_kernels.cu:461 route_kernel - softplus scores, their top k with bias, weights from the unbiased scores. One +// threadgroup (512) a token, as the CUDA kernel; rows 1024 values apart, the picks kPicks = 8 apart. +kernel void ds4_route(device const float* logits [[buffer(0)]], device const float* bias [[buffer(1)]], + constant uint& n [[buffer(2)]], constant uint& k [[buffer(3)]], constant float& route_scale [[buffer(4)]], + device int* idx [[buffer(5)]], device float* w [[buffer(6)]], device const float* bias_vl [[buffer(7)]], + constant uint& has_vl [[buffer(8)]], device const uint* vlw [[buffer(9)]], + uint tg [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort wp [[simdgroup_index_in_threadgroup]]) { + threadgroup float sc[1024], bs[1024]; + threadgroup float bv[32]; + threadgroup int bi[32]; + threadgroup int pick[16]; + const uint t = tg; + device const float* l = logits + (size_t) t * 1024u; + device const float* bp = bias; + if (has_vl != 0u && ((vlw[t >> 5u] >> (t & 31u)) & 1u) != 0u) bp = bias_vl; + for (uint i = tid; i < n; i += 512u) { + const float v = l[i]; + const float sp = v > 20.0f ? v : log(1.0f + exp(v)); + sc[i] = sqrt(sp); + bs[i] = sc[i] + bp[i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint r = 0; r < k; ++r) { + float best = -INFINITY; + int bix = 0x7fffffff; + for (uint i = tid; i < n; i += 512u) + if (bs[i] > best || (bs[i] == best && int(i) < bix)) { + best = bs[i]; + bix = int(i); + } + for (uint o = 16u; o > 0u; o >>= 1u) { + const float ov = simd_shuffle_xor(best, o); + const int oi = simd_shuffle_xor(bix, o); + if (ov > best || (ov == best && oi < bix)) { + best = ov; + bix = oi; + } + } + if (lane == 0u) { + bv[wp] = best; + bi[wp] = bix; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0u) { + float b = -INFINITY; + int ix = 0x7fffffff; + for (uint q = 0; q < 16u; ++q) + if (bv[q] > b || (bv[q] == b && bi[q] < ix)) { + b = bv[q]; + ix = bi[q]; + } + pick[r] = ix; + bs[ix] = -INFINITY; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0u) { + float s = 0.0f; + for (uint r = 0; r < k; ++r) s += sc[pick[r]]; + for (uint r = 0; r < k; ++r) { + idx[(size_t) t * 8u + r] = pick[r]; + w[(size_t) t * 8u + r] = sc[pick[r]] / (s + 1e-20f) * route_scale; + } + } +} + +// ds4_kernels.cu:1336 gather_kernel and :1342 combine_kernel - a token's gathered expert-output rows and the MoE's sum. +struct Ds4Gather { + int src[16 * 8]; +}; +struct Ds4Combine { + int n[16]; + int row[16][8]; + float w[16][8]; + int cpu[16]; +}; + +kernel void ds4_gather(device const float* x [[buffer(0)]], constant Ds4Gather& g [[buffer(1)]], + constant uint& cols [[buffer(2)]], device float* out [[buffer(3)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const uint r = tgp.y; + const uint c = tgp.x * 256u + tid; + if (c < cols) out[(size_t) r * cols + c] = x[(size_t) g.src[r] * cols + c]; +} + +kernel void ds4_combine(device const float* ed [[buffer(0)]], device const float* cpu [[buffer(1)]], + device const float* sh [[buffer(2)]], constant Ds4Combine& cm [[buffer(3)]], + constant uint& dim [[buffer(4)]], constant uint& has_cpu [[buffer(5)]], + constant uint& has_sh [[buffer(6)]], device float* out [[buffer(7)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const uint e = tgp.x * 256u + tid, t = tgp.y; + if (e >= dim) return; + float a = 0.0f; + for (int j = 0; j < cm.n[t]; ++j) a = fma(cm.w[t][j], ed[(size_t) cm.row[t][j] * dim + e], a); + if (has_cpu != 0u && cm.cpu[t] != 0) a += cpu[(size_t) t * dim + e]; + if (has_sh != 0u) a += sh[(size_t) t * dim + e]; + out[(size_t) t * dim + e] = a; +} + +// ---------------------------------------------------------------- ds4_kernels.cu:1156 markov_kernel / :1265 confidence_kernel +// +// The drafter's head (DSpark's Markov chain and its confidence): markov_add is a warp a vocabulary row - the previous +// token's E row (BF16) staged as floats in the threadgroup, then E . H[v] added to logits[v] (H is F16, the pack's +// markov_head.head.weight); confidence is a threadgroup a token over x and the same E row against the head's weights, +// its block sum through the logistic. +struct Ds4MarkovJob { + device float* logits; // [V]: += the bias + device const ushort* E; // [vocab][R] BF16 + device const ushort* H; // [vocab][R] F16 + device const int* tok; // the previous token (one int) + int V, R; +}; + +struct Ds4ConfJob { + device const float* x; // [T][dim] + device const ushort* E; // [vocab][R] BF16 + device const int* tok; // [T] + device const float* w; // [dim + R] + device float* out; // [T] + int dim, R; +}; + +// a BF16 bit pattern as its float (ds4_kernels.cu:28 bf16f: the top half, no rounding) +static inline float ds4_bf16f(ushort b) { return as_type(uint(b) << 16u); } + +// a warp a vocabulary row: logits[v] += E[*tok] . H[v] +kernel void ds4_markov_add(constant Ds4MarkovJob& j [[buffer(0)]], uint tg [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float er[1024]; // (R <= 1024: the host wrapper refuses more) + const device ushort* erow = j.E + (size_t) (*j.tok) * (size_t) j.R; + for (int i = int(tid); i < j.R; i += 256) er[i] = ds4_bf16f(erow[i]); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int v = int(tg) * 8 + int(tid >> 5); + if (v >= j.V) return; + const device ushort* hr = j.H + (size_t) v * (size_t) j.R; + float a = 0.0f; + for (int c = int(lane) * 8; c < j.R; c += 256) { + for (int i = 0; i < 8; ++i) a += float(as_type(hr[c + i])) * er[c + i]; + } + a = ds4_warp_sum(a); + if (lane == 0u) j.logits[v] += a; +} + +// a threadgroup a token: the logistic of the head's sum over x and the token's E row +kernel void ds4_confidence(constant Ds4ConfJob& j [[buffer(0)]], uint tg [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]], + ushort w [[simdgroup_index_in_threadgroup]]) { + threadgroup float red[32]; + const int t = int(tg); + const device float* xt = j.x + (size_t) t * (size_t) j.dim; + const device ushort* er = j.E + (size_t) j.tok[t] * (size_t) j.R; + float a = 0.0f; + for (int i = int(tid); i < j.dim + j.R; i += 1024) a += j.w[i] * (i < j.dim ? xt[i] : ds4_bf16f(er[i - j.dim])); + a = ds4_block_sum_n(a, red, lane, w, 32u); + if (tid == 0u) j.out[t] = 1.0f / (1.0f + exp(-a)); +} + +// embed by raw device addresses (the runtime's MTLResidencySet makes gpuAddress pointers valid kernel arguments) +struct Ds4EmbedArgs { + device ushort* rows; + device float* h; + uint dim; + uint hc; +}; +kernel void ds4_embed_j(device const Ds4EmbedArgs* a [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + ushort lane [[thread_index_in_threadgroup]]) { + const uint e = tgp.x * 256u + uint(lane); + if (e >= a->dim) return; + const uint t = tgp.y; + const uint bits = uint(a->rows[(size_t) t * a->dim + e]) << 16; + const float v = as_type(bits); + for (uint s = 0; s < a->hc; ++s) a->h[((size_t) t * a->hc + s) * a->dim + e] = v; +} + +// (diagnostic) write a constant through a device pointer passed as a value, and one passed in a buffer +kernel void ds4_test_write(device float* out [[buffer(0)]], constant uint& n [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + if (gid < n) out[gid] = 3.0f; +} +kernel void ds4_test_write_v(device const uint64_t* p [[buffer(0)]], constant uint& n [[buffer(1)]], uint gid [[thread_position_in_grid]]) { + device float* out = (device float*) (uintptr_t) p[0]; + if (gid < n) out[gid] = 7.0f; +} + +// (diagnostic) report the arguments the kernel received +kernel void ds4_args_probe(device float* out [[buffer(0)]], constant uint& dim [[buffer(1)]], + constant uint& hc [[buffer(2)]], uint gid [[thread_position_in_grid]]) { + if (gid == 0) { out[0] = float(dim); out[1] = float(hc); out[2] = 1.0f; } +} + +// hc == 4 (every layer of this model): the warp-lane form of the CUDA kernel - lane j*4+k holds comb[j][k], the row and column +// sums are xor-shuffles (1, 2 within a row; 4, 8 within a column). One simdgroup a token instead of one lane. +kernel void ds4_hc_split4(device const float* mixes [[buffer(0)]], device const float* scale [[buffer(1)]], + device const float* base [[buffer(2)]], device float* pre [[buffer(3)]], + device float* post [[buffer(4)]], device float* comb [[buffer(5)]], + constant uint& hc [[buffer(6)]], constant uint& iters [[buffer(7)]], + constant float& eps [[buffer(8)]], uint b [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_threadgroup]]) { + const device float* mx = mixes + (size_t) b * 32u; + if (lane < 4u) { + pre[(size_t) b * 8u + lane] = 1.0f / (1.0f + exp(-(mx[lane] * scale[0] + base[lane]))) + eps; + post[(size_t) b * 8u + lane] = 2.0f / (1.0f + exp(-(mx[4u + lane] * scale[1] + base[4u + lane]))); + } + const bool on = lane < 16u; + float c = on ? mx[8u + lane] * scale[2] + base[8u + lane] : -3.402823466e+38f; + float m = fmax(c, simd_shuffle_xor(c, ushort(1))); + m = fmax(m, simd_shuffle_xor(m, ushort(2))); + c = on ? exp(c - m) : 0.0f; + float rs = c + simd_shuffle_xor(c, ushort(1)); + rs += simd_shuffle_xor(rs, ushort(2)); + c = c / rs + eps; + c = on ? c : 0.0f; + float cs = c + simd_shuffle_xor(c, ushort(4)); + cs += simd_shuffle_xor(cs, ushort(8)); + c = c / (cs + eps); + for (uint it = 1; it < iters; ++it) { + float r = c + simd_shuffle_xor(c, ushort(1)); + r += simd_shuffle_xor(r, ushort(2)); + c = c / (r + eps); + float s = c + simd_shuffle_xor(c, ushort(4)); + s += simd_shuffle_xor(s, ushort(8)); + c = c / (s + eps); + } + if (on) comb[(size_t) b * 64u + lane] = c; +} + +// Plain byte copies and fills on the context's queue: cudaMemcpyAsync / cudaMemsetAsync between runtime allocations (unified +// memory: a "device" and a "host" buffer are both just memory) ride in the same command buffer as the kernels around them, +// so their order is the program order and a layer needs no queue-to-queue handshake. +kernel void ds4_copy4(device const uint* src [[buffer(0)]], device uint* dst [[buffer(1)]], constant ulong& n4 [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (ulong(gid) < n4) dst[gid] = src[gid]; +} +kernel void ds4_copy1(device const uchar* src [[buffer(0)]], device uchar* dst [[buffer(1)]], constant ulong& n [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (ulong(gid) < n) dst[gid] = src[gid]; +} +kernel void ds4_fill1(device uchar* dst [[buffer(0)]], constant ulong& n [[buffer(1)]], constant uint& v [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + if (ulong(gid) < n) dst[gid] = uchar(v); +} + +// --------------------------------------------------------------------------------------------------------------------- +// The prompt chunk's dense GEMM operands (see ds4_gemm_f16_nt on the host): float rows -> half, a quantized matrix -> half rows, +// strided half rows packed, and the float result scattered into the engine's strided output. + +kernel void ds4_rows_f16(device const float* x [[buffer(0)]], constant ulong& ldx [[buffer(1)]], + device const int* idx [[buffer(2)]], constant uint& has_idx [[buffer(3)]], + constant uint& ncols [[buffer(4)]], device half* out [[buffer(5)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const uint r = tgp.y, c = tgp.x * 256u + tid; + if (c >= ncols) return; + const ulong row = has_idx != 0u ? ulong(idx[r]) : ulong(r); + out[(ulong) r * ncols + c] = half(x[row * ldx + c]); +} + +static inline void ds4_kscale2(const device uchar* q, int j, thread int& d, thread int& m) { + if (j < 4) { + d = int(q[j]) & 63; + m = int(q[j + 4]) & 63; + } else { + d = (int(q[j + 4]) & 0xf) | ((int(q[j - 4]) >> 6) << 4); + m = (int(q[j + 4]) >> 4) | ((int(q[j]) >> 6) << 4); + } +} + +// a thread a 32-value sub-block of a Q5_K matrix (n rows of k values, ld bytes a row) -> out[n][k] half +kernel void ds4_dq_q5k(device const uchar* W [[buffer(0)]], constant ulong& ld [[buffer(1)]], constant uint& k [[buffer(2)]], + constant uint& n [[buffer(3)]], device half* out [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + const uint nb = k / 256u; + const ulong total = (ulong) n * nb * 8ul; + if (ulong(gid) >= total) return; + const uint ib = gid & 7u; + const uint blk = (gid >> 3) % nb; + const uint row = uint(ulong(gid) / (8ul * nb)); + const device uchar* p = W + (ulong) row * ld + (ulong) blk * 176ul; + const float d = float(as_type(*(const device ushort*) p)); + const float dmin = float(as_type(*(const device ushort*) (p + 2))); + int sc, mm; + ds4_kscale2(p + 4, int(ib), sc, mm); + const device uchar* qh = p + 16; + const device uchar* qs = p + 48 + 32u * (ib >> 1); + const uint sh = 4u * (ib & 1u); + device packed_half4* o = (device packed_half4*) (out + (ulong) row * k + (ulong) blk * 256ul + ib * 32u); + const float ds = d * float(sc), dm = dmin * float(mm); + for (uint l = 0; l < 32u; l += 4u) { // four values a step: 4-byte loads, 8-byte stores (the same arithmetic a value) + const uint4 a = uint4(*(const device packed_uchar4*) (qs + l)); + const uint4 b = uint4(*(const device packed_uchar4*) (qh + l)); + const uint4 q = ((a >> sh) & 15u) | (((b >> ib) & 1u) << 4); + o[l >> 2] = packed_half4(half4(float4(ds) * float4(q) - float4(dm))); + } +} + +// Q6_K (210-byte blocks: ql[128], qh[64], int8 scales[16], f16 d; the Unsloth sets' attention, shared-expert, head and engram_kv +// matrices): a thread four values of a block (two 128-value halves x 32 lanes), ggml's dequantize_row_q6_K arithmetic +// ((d * scale) * (q - 32), in this order) -> out[n][k] half +kernel void ds4_dq_q6k(device const uchar* W [[buffer(0)]], constant ulong& ld [[buffer(1)]], constant uint& k [[buffer(2)]], + constant uint& n [[buffer(3)]], device half* out [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + const uint nb = k / 256u; + const ulong total = (ulong) n * nb * 64ul; + if (ulong(gid) >= total) return; + const uint l = gid & 31u; + const uint hf = (gid >> 5) & 1u; + const uint blk = (gid >> 6) % nb; + const uint row = uint(ulong(gid) / (64ul * nb)); + const device uchar* p = W + (ulong) row * ld + (ulong) blk * 210ul; + const float d = float(as_type(*(const device ushort*) (p + 208))); + const device uchar* ql = p + 64u * hf; + const device uchar* qh = p + 128u + 32u * hf; + const device char* sc = (const device char*) (p + 192u + 8u * hf); + const uint is = l >> 4; + const uint qlo = uint(ql[l]), qhi = uint(ql[l + 32u]), h = uint(qh[l]); + const int q1 = int((qlo & 0xfu) | (((h >> 0) & 3u) << 4)) - 32; + const int q2 = int((qhi & 0xfu) | (((h >> 2) & 3u) << 4)) - 32; + const int q3 = int((qlo >> 4) | (((h >> 4) & 3u) << 4)) - 32; + const int q4 = int((qhi >> 4) | (((h >> 6) & 3u) << 4)) - 32; + device half* o = out + (ulong) row * k + (ulong) blk * 256ul + 128ul * hf + l; + o[0] = half(d * float(sc[is + 0u]) * float(q1)); + o[32] = half(d * float(sc[is + 2u]) * float(q2)); + o[64] = half(d * float(sc[is + 4u]) * float(q3)); + o[96] = half(d * float(sc[is + 6u]) * float(q4)); +} + +// Q8_0: 34-byte blocks (fp16 d, 32 int8) +kernel void ds4_dq_q8(device const uchar* W [[buffer(0)]], constant ulong& ld [[buffer(1)]], constant uint& k [[buffer(2)]], + constant uint& n [[buffer(3)]], device half* out [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + const uint ng = k / 32u; + const ulong total = (ulong) n * ng; + if (ulong(gid) >= total) return; + const uint g = gid % ng, row = gid / ng; + const device uchar* p = W + (ulong) row * ld + (ulong) g * 34ul; + const float d = float(as_type(*(const device ushort*) p)); + device half* o = out + (ulong) row * k + (ulong) g * 32ul; + for (uint l = 0; l < 32u; ++l) o[l] = half(d * float(char(p[2 + l]))); +} + +// BF16 values (ld values a row) -> half +kernel void ds4_dq_bf16(device const ushort* W [[buffer(0)]], constant ulong& ld [[buffer(1)]], constant uint& k [[buffer(2)]], + constant uint& n [[buffer(3)]], device half* out [[buffer(4)]], uint gid [[thread_position_in_grid]]) { + const ulong total = (ulong) n * k; + if (ulong(gid) >= total) return; + const ulong r = ulong(gid) / k, c = ulong(gid) - r * k; + out[gid] = half(as_type(uint(W[r * ld + c]) << 16)); +} + +// half rows with a row stride -> contiguous rows +kernel void ds4_pack_h16(device const ushort* x [[buffer(0)]], constant ulong& ld [[buffer(1)]], constant uint& k [[buffer(2)]], + device ushort* out [[buffer(3)]], uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const uint c = tgp.x * 256u + tid; + if (c < k) out[(ulong) tgp.y * k + c] = x[(ulong) tgp.y * ld + c]; +} + +// y[r][c] = alpha * t[r][c] + beta * y[r][c] +kernel void ds4_axpby_rows(device const float* t [[buffer(0)]], device float* y [[buffer(1)]], constant ulong& ldy [[buffer(2)]], + constant uint& n [[buffer(3)]], constant float& alpha [[buffer(4)]], constant float& beta [[buffer(5)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const uint c = tgp.x * 256u + tid; + if (c >= n) return; + const ulong yi = (ulong) tgp.y * ldy + c; + y[yi] = alpha * t[(ulong) tgp.y * n + c] + (beta != 0.0f ? beta * y[yi] : 0.0f); +} + +// ds4_q8_quant for several jobs in one dispatch (a matvec call's activations): grid (groups, max rows, jobs) +struct Ds4Q8Spec { + device const float* x; + ulong ldx; + uint rows; + uint k; + device char* q; + device float* d; +}; +kernel void ds4_q8_quant_multi(device const Ds4Q8Spec* specs [[buffer(0)]], uint3 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + const device Ds4Q8Spec& s = specs[tgp.z]; + const uint warp = tid / 32u, lane = tid % 32u; + const uint b = tgp.x * 8u + warp; + const uint r = tgp.y; + if (r >= s.rows || (size_t) b * 32u >= (size_t) s.k) return; + const device float* xr = s.x + (size_t) r * s.ldx; + const float v = xr[(size_t) b * 32u + lane]; + const float am = simd_max(fabs(v)); + const float sc = am / 127.0f; + s.q[(size_t) r * s.k + (size_t) b * 32u + lane] = char(am > 0.0f ? rint(v / sc) : 0.0f); + if (lane == 0u) s.d[(size_t) r * (s.k / 32u) + b] = sc; +} diff --git a/edge0-strata/metal/src/engine/msl/ds4_iq.metal b/edge0-strata/metal/src/engine/msl/ds4_iq.metal new file mode 100644 index 00000000..5bf41e04 --- /dev/null +++ b/edge0-strata/metal/src/engine/msl/ds4_iq.metal @@ -0,0 +1,1184 @@ +#include +using namespace metal; +// ds4_iq.cu:809 mv_kernel - the I-quant expert matvec. Same frame as the MXFP4/Q8_0 ones (a warp a row, a lane strides the +// 256-value blocks, sixteen accumulator columns for the token window), with the block decoder chosen per job by b.type: +// 16 IQ2_XXS (66 B), 17 IQ2_XS (74 B), 22 IQ2_S (82 B). The activations are the quantiser's int8 (xq) with their scales (xd), +// so each sub-group's dot is integer and the two scales fold in at the end - the same arithmetic as the float probes, at a +// quarter of the traffic. +struct Ds4MvJob { + device const uchar* w; + ulong ld; + int type, k, n, ncol; + device const char* xq; + device const float* xd; + device float* y; + ulong ycs; + device const float* bias; + float alpha; +}; + +kernel void ds4_iq_mv(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nsb = uint(b.k) / 32u; + const uint nblk = uint(b.k) / 256u; + const uint bs = b.type == 16 ? 66u : (b.type == 17 ? 74u : (b.type == 18 ? 98u : 82u)); + const device uchar* w = b.w + (size_t) row * b.ld; + float acc[16]; + for (uint c = 0; c < 16u; ++c) acc[c] = 0.0f; + for (uint blk = lane; blk < nblk; blk += 32u) { + const device uchar* p = w + (size_t) blk * bs; + const float d = float(as_type(*(const device ushort*) p)); + for (uint ib = 0; ib < 8u; ++ib) { + if (b.type == 18) { // IQ3_XXS: d, qs[64] (two grid bytes a 8 values), scales_and_signs[8] words (top nibble the scale) + const device uchar* qs = p + 2; + const uint w0 = uint(qs[8 * ib]) | (uint(qs[8 * ib + 1]) << 8) | (uint(qs[8 * ib + 2]) << 16) | (uint(qs[8 * ib + 3]) << 24); + const uint w1 = uint(qs[8 * ib + 4]) | (uint(qs[8 * ib + 5]) << 8) | (uint(qs[8 * ib + 6]) << 16) | (uint(qs[8 * ib + 7]) << 24); + const uint aux = uint(qs[64 + 4 * ib]) | (uint(qs[65 + 4 * ib]) << 8) | (uint(qs[66 + 4 * ib]) << 16) | (uint(qs[67 + 4 * ib]) << 24); + const float db3 = d * (0.5f + float(aux >> 28)) * 0.5f; + for (uint l = 0; l < 4u; ++l) { + const uint wv = l < 2u ? w0 : w1, sh = 16u * (l & 1u); + const uint sgn = uint(ksigns_iq2xs[(aux >> (7u * l)) & 127u]); + const uint g0 = iq3xxs_grid[(wv >> sh) & 255u], g1 = iq3xxs_grid[(wv >> (sh + 8u)) & 255u]; + const uint off = (size_t) blk * 256u + ib * 32u + l * 8u; + for (uint c = 0; c < 16u; ++c) { + if (int(c) >= b.ncol) break; + const device char* xq = b.xq + (size_t) c * b.k + off; + int sum = 0; + for (uint j = 0; j < 4u; ++j) { + int v0 = int((g0 >> (8u * j)) & 0xffu), v1 = int((g1 >> (8u * j)) & 0xffu); + if ((sgn & (1u << j)) != 0u) v0 = -v0; + if ((sgn & (1u << (j + 4u))) != 0u) v1 = -v1; + sum += v0 * int(xq[j]) + v1 * int(xq[j + 4u]); + } + acc[c] = fma(b.xd[(size_t) c * nsb + (off / 32u)] * db3, float(sum), acc[c]); + } + } + continue; + } + float db0, db1; + uint idx[4]; + uint sg[4]; + if (b.type == 16) { + const uint aux0 = uint(p[2 + 8 * ib]) | (uint(p[3 + 8 * ib]) << 8) | (uint(p[4 + 8 * ib]) << 16) | + (uint(p[5 + 8 * ib]) << 24); // (8 bytes a sub-block: grid bytes, then signs and scale) + const uint aux1 = uint(p[6 + 8 * ib]) | (uint(p[7 + 8 * ib]) << 8) | (uint(p[8 + 8 * ib]) << 16) | + (uint(p[9 + 8 * ib]) << 24); + const float db = d * (0.5f + float(aux1 >> 28)) * 0.25f; + db0 = db1 = db; + for (uint l = 0; l < 4u; ++l) { + idx[l] = (aux0 >> (8u * l)) & 0xffu; + sg[l] = uint(ksigns_iq2xs[(aux1 >> (7u * l)) & 127u]); + } + } else if (b.type == 17) { + const device ushort* qs = (const device ushort*) (p + 2); + const device uchar* sc = p + 66; + db0 = d * (0.5f + float(sc[ib] & 0xfu)) * 0.25f; + db1 = d * (0.5f + float(sc[ib] >> 4)) * 0.25f; + for (uint l = 0; l < 4u; ++l) { + const uint q = uint(qs[4u * ib + l]); + idx[l] = q & 511u; + sg[l] = uint(ksigns_iq2xs[(q >> 9) & 127u]); // the sign byte of the 7 stored bits (bit j set: value j negated) + } + } else { + const device uchar* qs = p + 2; + const device uchar* sgb = qs + 32; + const device uchar* qh = p + 66; + const device uchar* sc = p + 74; + db0 = d * (0.5f + float(sc[ib] & 0xfu)) * 0.25f; + db1 = d * (0.5f + float(sc[ib] >> 4)) * 0.25f; + for (uint l = 0; l < 4u; ++l) { + idx[l] = uint(qs[4u * ib + l]) | ((uint(qh[ib]) << (8u - 2u * l)) & 0x300u); + sg[l] = uint(sgb[4u * ib + l]); + } + } + for (uint l = 0; l < 4u; ++l) { + ulong g; + if (b.type == 16) g = iq2xxs_grid[idx[l]]; + else if (b.type == 17) g = iq2xs_grid[idx[l]]; + else g = iq2s_grid[idx[l] & 1023u]; + const float db = (l < 2u) ? db0 : db1; + const uint off = (size_t) blk * 256u + ib * 32u + l * 8u; + for (uint c = 0; c < 16u; ++c) { + if (int(c) >= b.ncol) break; + const device char* xq = b.xq + (size_t) c * b.k + off; + int sum = 0; + for (uint j = 0; j < 8u; ++j) { + const int code = int((g >> (8u * j)) & 0xffu); + const bool neg = b.type == 16 ? ((sg[l] & uint(kmask_iq2xs[j])) != 0u) + : ((sg[l] & (1u << j)) != 0u); + sum += (neg ? -code : code) * int(xq[j]); + } + acc[c] = fma(b.xd[(size_t) c * nsb + (off / 32u)] * db, float(sum), acc[c]); + } + } + } + } + for (uint c = 0; c < 16u; ++c) { + if (int(c) >= b.ncol) break; + float v = acc[c]; + v += simd_shuffle_xor(v, 16); + v += simd_shuffle_xor(v, 8); + v += simd_shuffle_xor(v, 4); + v += simd_shuffle_xor(v, 2); + v += simd_shuffle_xor(v, 1); + if (lane == 0u) b.y[(size_t) c * b.ycs + row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); + } +} + +// Q5_K / Q6_K (the dense matrices' K-quants; ggml blocks of 176 / 210 bytes): the same warp a row frame, a lane a +// 256-value block, each value's code times its activation's int8. Q5_K's weight is d*sc*q - dmin*m (both parts of the +// 32-value group's scales), Q6_K's d*sc*(q - 32). +static inline void ds4_k_scale(const device uchar* q, int j, thread int& d, thread int& m) { + if (j < 4) { + d = int(q[j]) & 63; + m = int(q[j + 4]) & 63; + } else { + d = (int(q[j + 4]) & 0xf) | ((int(q[j - 4]) >> 6) << 4); + m = (int(q[j + 4]) >> 4) | ((int(q[j]) >> 6) << 4); + } +} + +kernel void ds4_iq_mv_k(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nsb = uint(b.k) / 32u; + const uint nblk = uint(b.k) / 256u; + const uint bs = b.type == 13 ? 176u : 210u; + const device uchar* w = b.w + (size_t) row * b.ld; + float acc[16]; + for (uint c = 0; c < 16u; ++c) acc[c] = 0.0f; + for (uint blk = lane; blk < nblk; blk += 32u) { + const device uchar* p = w + (size_t) blk * bs; + if (b.type == 13) { + const float d = float(as_type(*(const device ushort*) p)); + const float dmin = float(as_type(*(const device ushort*) (p + 2))); + const device uchar* scales = p + 4; + const device uchar* qh = p + 16; + const device uchar* qs = p + 48; + for (uint ib = 0; ib < 8u; ++ib) { + int sc, mm; + ds4_k_scale(scales, int(ib), sc, mm); + const device uchar* qsp = qs + 32u * (ib >> 1); + const uint sh = 4u * (ib & 1u); + for (uint t = 0; t < 8u; ++t) { + const uint q4 = (*(const device uint*) (qsp + 4u * t) >> sh) & 0x0f0f0f0fu; + const uint h4 = (*(const device uint*) (qh + 4u * t) >> ib) & 0x01010101u; + const uint q = q4 | (h4 << 4); + const uint off = blk * 256u + ib * 32u + t * 4u; + for (uint c = 0; c < 16u; ++c) { + if (int(c) >= b.ncol) break; + const device char* xq = b.xq + (size_t) c * b.k + off; + int sumq = 0, sumx = 0; + for (uint j = 0; j < 4u; ++j) { + const int code = int((q >> (8u * j)) & 0xffu); + const int xv = int(xq[j]); + sumq += code * xv; + sumx += xv; + } + const float dot = d * float(sc) * float(sumq) - dmin * float(mm) * float(sumx); + acc[c] = fma(b.xd[(size_t) c * nsb + (off / 32u)], dot, acc[c]); + } + } + } + } else { + const float d = float(as_type(*(const device ushort*) (p + 208))); + const device uchar* ql = p; + const device uchar* qh = p + 128; + const device uchar* scl = p + 192; + for (uint ib = 0; ib < 8u; ++ib) { + const uint h = ib >> 2, quarter = ib & 3u; + const device uchar* qlp = ql + 64u * h + ((quarter & 1u) ? 32u : 0u); + const device uchar* qhp = qh + 32u * h; + const uint sh = 4u * (quarter >> 1); + const int s0 = int(((const device char*) scl)[8u * h + 2u * quarter]); + const int s1 = int(((const device char*) scl)[8u * h + 2u * quarter + 1u]); + for (uint t = 0; t < 8u; ++t) { + // (packed_uchar4: a 210-byte block starts only 2-byte aligned - a plain uint load reads the wrong word on odd blocks) + const uint x = ((as_type(uchar4(*(const device packed_uchar4*) (qlp + 4u * t))) >> sh) & 0x0f0f0f0fu) | + (((as_type(uchar4(*(const device packed_uchar4*) (qhp + 4u * t))) >> (2u * quarter)) & 0x03030303u) << 4); + const uint off = blk * 256u + ib * 32u + t * 4u; + const float s = float(t < 4u ? s0 : s1); + for (uint c = 0; c < 16u; ++c) { + if (int(c) >= b.ncol) break; + const device char* xq = b.xq + (size_t) c * b.k + off; + int sumq = 0; + for (uint j = 0; j < 4u; ++j) { + const int code = int((x >> (8u * j)) & 0xffu) - 32; + sumq += code * int(xq[j]); + } + acc[c] = fma(b.xd[(size_t) c * nsb + (off / 32u)] * d * s, float(sumq), acc[c]); + } + } + } + } + } + for (uint c = 0; c < 16u; ++c) { + if (int(c) >= b.ncol) break; + float v = acc[c]; + v += simd_shuffle_xor(v, 16); + v += simd_shuffle_xor(v, 8); + v += simd_shuffle_xor(v, 4); + v += simd_shuffle_xor(v, 2); + v += simd_shuffle_xor(v, 1); + if (lane == 0u) b.y[(size_t) c * b.ycs + row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); + } +} + +// ---------------------------------------------------------------- ds4_iq.cu:536 gate_up_kernel / :588 down_kernel +// +// The prompt / window path's expert GEMM (ds4_iq.cu's launch_gu / launch_dn): a threadgroup of 256 = 8 warps, a warp +// RW rows of the gate and of the up (down: RW rows of the down), a lane 32 / RW sub-blocks of the row, and the group's +// entries NE at a pass so that a sub-block the weights decode once serves them all. Grid (rows / (8 RW), groups). +// +// The group's blob and its entries are the CUDA kernel's `group` (ds4_iq.cu:512): a plan group (grp_ptr null) is one +// expert whose entry is the group's own index and whose activation row is 0, a window group's entries come from +// grp_ent[grp_start[g] .. grp_start[g + 1]) and an entry e's activation row is e >> 3. +struct Ds4IqJob { + device const ulong* plan_ptr; // plan_ptr[g]: the expert's blob (0: the group is skipped) + device const ulong* grp_ptr; // grp_ptr[g]: a window set's blob + device const int* grp_start; // a set's first entry (n_groups + 1 of them) + device const int* grp_ent; // the entries (token * 8 + pick) + device const int* n_groups; + device const char* xq; // the int8 activations of a row (n_embd an entry) + device const float* xd; // their scales (n_embd / 32 a row) + device const char* hq; // down's int8 input (n_ff an entry) + device const float* hd; + device float* h; // [entries][n_ff]: the clamped swiglu + device float* part; // [entries][n_embd]: each entry's down output + long tr0, tr1, tr2; // the roles' offsets in a blob + int n_embd, n_ff; + int type; // 16 IQ2_XXS, 17 IQ2_XS, 22 IQ2_S (ds4_iq.cuh's kIQ2_*) + float limit; // swiglu_limit (0: none) +}; + +// a 32-value sub-block as 32 int8 codes (8 packed ints, values 4 w .. 4 w + 3 in word w) and the two halves' scales +// (ggml's d (2 s + 1) / 8, the .cu's Sub.d * Sub.s0 in one F32: the codes of the second half are scaled by s1) +struct Ds4IqSub { + int q[8]; + float s0, s1; +}; + +// the block bytes of the types these kernels take: 16 IQ2_XXS, 17 IQ2_XS, 18 IQ3_XXS, 22 IQ2_S (ds4_iq.cu's Ty::B) +static inline uint ds4_iq_bs(int type) { + return type == 16 ? 66u : type == 17 ? 74u : type == 18 ? 98u : type == 19 ? 50u : type == 29 ? 56u : 82u; +} + +// the 1-bit types' sub-block decoders (defined with the multi-column matvec below): 32 codes as 8*grid +- 1 (int8) and the two halves' scales +static inline void ds4_w32_iq1s(const device uchar* p, uint ib, thread int (&w)[32], thread float& db0, thread float& db1); +static inline void ds4_w32_iq1m(const device uchar* p, uint ib, thread int (&w)[32], thread float& db0, thread float& db1); + +// ds4_iq.cu:139 decode: the block's 8 sub-blocks, kb's is the (kb & 7)'th; the grids are ggml-common.h's +static inline Ds4IqSub ds4_iq_sub(device const uchar* row, uint bs, int kb, int type) { + Ds4IqSub s; + const int ib = kb & 7; + const device uchar* p = row + (size_t) (kb >> 3) * bs; + const float d = float(as_type(*(const device ushort*) p)); + uint idx[4], sm[4]; + float db0, db1; + if (type == 19 || type == 29) { // IQ1_S (50 B) / IQ1_M (56 B): the integer form 8*code +- 1, the scale / 8 in db + int w[32]; + if (type == 19) ds4_w32_iq1s(p, uint(ib), w, db0, db1); else ds4_w32_iq1m(p, uint(ib), w, db0, db1); + for (uint t = 0; t < 8u; ++t) + s.q[t] = (w[4u * t] & 0xff) | ((w[4u * t + 1u] & 0xff) << 8) | ((w[4u * t + 2u] & 0xff) << 16) | ((w[4u * t + 3u] & 0xff) << 24); + s.s0 = db0; + s.s1 = db1; + return s; + } + if (type == 18) { // IQ3_XXS: d, qs[64] grid bytes (4 a sub-block), scales[8] (the top nibble the group's s) + const device uchar* qs = p + 2; + const uint w0 = uint(qs[8 * ib]) | (uint(qs[8 * ib + 1]) << 8) | (uint(qs[8 * ib + 2]) << 16) | + (uint(qs[8 * ib + 3]) << 24); + const uint w1 = uint(qs[8 * ib + 4]) | (uint(qs[8 * ib + 5]) << 8) | (uint(qs[8 * ib + 6]) << 16) | + (uint(qs[8 * ib + 7]) << 24); + const uint aux = uint(qs[64 + 4 * ib]) | (uint(qs[65 + 4 * ib]) << 8) | (uint(qs[66 + 4 * ib]) << 16) | + (uint(qs[67 + 4 * ib]) << 24); + for (uint l = 0; l < 4u; ++l) { + const uint w = l < 2u ? w0 : w1, sh = 16u * (l & 1u); + const uint sg = uint(ksigns_iq2xs[(aux >> (7u * l)) & 127u]); + const uint g0 = iq3xxs_grid[(w >> sh) & 255u], g1 = iq3xxs_grid[(w >> (sh + 8u)) & 255u]; + int a0 = 0, a1 = 0; + for (uint j = 0; j < 4u; ++j) { + int v0 = int((g0 >> (8u * j)) & 0xffu), v1 = int((g1 >> (8u * j)) & 0xffu); + if ((sg & (1u << j)) != 0u) v0 = -v0; + if ((sg & (1u << (j + 4u))) != 0u) v1 = -v1; + a0 |= (v0 & 0xff) << int(8u * j); + a1 |= (v1 & 0xff) << int(8u * j); + } + s.q[2u * l] = a0; + s.q[2u * l + 1u] = a1; + } + const float ds = d * (0.5f + float(aux >> 28)) * 0.5f; // ggml's db: d (0.5 + s) 0.5 + s.s0 = ds; + s.s1 = ds; + return s; + } + if (type == 16) { // d, qs[8] i16 (4 grid bytes, then 4 sign/scale bytes), scales in the top nibbles + const uint aux0 = uint(p[2 + 8 * ib]) | (uint(p[3 + 8 * ib]) << 8) | (uint(p[4 + 8 * ib]) << 16) | + (uint(p[5 + 8 * ib]) << 24); + const uint aux1 = uint(p[6 + 8 * ib]) | (uint(p[7 + 8 * ib]) << 8) | (uint(p[8 + 8 * ib]) << 16) | + (uint(p[9 + 8 * ib]) << 24); + const float db = d * (0.5f + float(aux1 >> 28)) * 0.25f; + db0 = db1 = db; + for (uint l = 0; l < 4u; ++l) { + idx[l] = (aux0 >> (8u * l)) & 0xffu; + sm[l] = uint(ksigns_iq2xs[(aux1 >> (7u * l)) & 127u]); + } + } else if (type == 17) { // d, qs[16] i16 (9-bit grid, 7 sign bits), scales[8] (two 4-bit) + const device ushort* qs = (const device ushort*) (p + 2); + const device uchar* sc = p + 66; + db0 = d * (0.5f + float(sc[ib] & 0x0fu)) * 0.25f; + db1 = d * (0.5f + float(sc[ib] >> 4)) * 0.25f; + for (uint l = 0; l < 4u; ++l) { + const uint q = uint(qs[4u * uint(ib) + l]); + idx[l] = q & 511u; + sm[l] = uint(ksigns_iq2xs[(q >> 9) & 127u]); + } + } else { // IQ2_S: d, qs[32] (8-bit grid), signs[32], qh[8] (2 bits more), scales[8] + const device uchar* qs = p + 2; + const device uchar* sgb = qs + 32; + const device uchar* qh = p + 66; + const device uchar* sc = p + 74; + db0 = d * (0.5f + float(sc[ib] & 0x0fu)) * 0.25f; + db1 = d * (0.5f + float(sc[ib] >> 4)) * 0.25f; + for (uint l = 0; l < 4u; ++l) { + idx[l] = uint(qs[4u * uint(ib) + l]) | ((uint(qh[ib]) << (8u - 2u * l)) & 0x300u); + sm[l] = uint(sgb[4u * uint(ib) + l]); + } + } + for (uint l = 0; l < 4u; ++l) { + const ulong g = type == 16 ? iq2xxs_grid[idx[l]] : (type == 17 ? iq2xs_grid[idx[l]] : iq2s_grid[idx[l] & 1023u]); + int w0 = 0, w1 = 0; + for (uint j = 0; j < 4u; ++j) { + int v0 = int((g >> (8u * j)) & 0xffu), v1 = int((g >> (8u * (j + 4u))) & 0xffu); + if ((sm[l] & (1u << j)) != 0u) v0 = -v0; + if ((sm[l] & (1u << (j + 4u))) != 0u) v1 = -v1; + w0 |= (v0 & 0xff) << int(8u * j); + w1 |= (v1 & 0xff) << int(8u * j); + } + s.q[2u * l] = w0; + s.q[2u * l + 1u] = w1; + } + s.s0 = db0; + s.s1 = db1; + return s; +} + +// 4 int8 codes against the 4 int8 activations packed in x (ds4_iq.cu:350's __dp4a: Metal has no such instruction) +static inline int ds4_iq_dot4(int q, int x) { + const char4 a = as_type(q), b = as_type(x); + return int(a.x) * int(b.x) + int(a.y) * int(b.y) + int(a.z) * int(b.z) + int(a.w) * int(b.w); +} + +// ds4_iq.cu:349 sub_dot: the sub-block against its 32 int8 activations (8 packed ints), the scales in F32 +static inline float ds4_iq_dot(Ds4IqSub s, device const int* xp) { + const int lo = ds4_iq_dot4(s.q[0], xp[0]) + ds4_iq_dot4(s.q[1], xp[1]) + ds4_iq_dot4(s.q[2], xp[2]) + + ds4_iq_dot4(s.q[3], xp[3]); + const int hi = ds4_iq_dot4(s.q[4], xp[4]) + ds4_iq_dot4(s.q[5], xp[5]) + ds4_iq_dot4(s.q[6], xp[6]) + + ds4_iq_dot4(s.q[7], xp[7]); + return s.s0 * float(lo) + s.s1 * float(hi); +} + +// ds4_iq.cu:502 lanes_sum: the row's sum over the 32 / RW lanes that share it +template +static inline float ds4_iq_lsum(float v) { + for (int o = 16 / RW; o > 0; o >>= 1) v += simd_shuffle_xor(v, ushort(o)); + return v; +} + +// ds4_iq.cu:497 swiglu: the limit's clamps, then g / (1 + e^-g) * u +static inline float ds4_iq_swiglu(float gv, float uv, float limit) { + if (limit > 0.0f) { + gv = min(gv, limit); + uv = min(max(uv, -limit), limit); + } + return (gv / (1.0f + exp(-gv))) * uv; +} + +// the group's blob and its entry range (ds4_iq.cu:512 group); false: the threadgroup has nothing to do +static inline bool ds4_iq_group(constant Ds4IqJob& j, uint g, thread ulong& blob, thread int& m0, thread int& m1) { + if (j.grp_ptr == nullptr) { + const ulong p = j.plan_ptr[g]; + if (p == 0ul) return false; + blob = p; + m0 = int(g); + m1 = m0 + 1; + return true; + } + if (int(g) >= *j.n_groups) return false; + blob = j.grp_ptr[g]; + m0 = j.grp_start[g]; + m1 = j.grp_start[g + 1]; + return true; +} + +// gate_up_kernel (kGuRW = 1: a warp a row, a lane a sub-block): h = the clamped swiglu of the gate and the up +kernel void ds4_iq_gu(constant Ds4IqJob& j [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + ulong blob; + int m0, m1; + if (!ds4_iq_group(j, tgp.y, blob, m0, m1)) return; + const int row = int(tgp.x) * 8 + int(tid >> 5); + if (row >= j.n_ff) return; + const uint bs = ds4_iq_bs(j.type); + const uint nsb = uint(j.n_embd) / 32u; + const size_t rb = (size_t) (j.n_embd / 256) * (size_t) bs; + const device uchar* g = (const device uchar*) blob + (size_t) j.tr0 + (size_t) row * rb; + const device uchar* u = (const device uchar*) blob + (size_t) j.tr1 + (size_t) row * rb; + for (int mm = m0; mm < m1; mm += 4) { + const int ne = min(4, m1 - mm); + int es[4], xr[4]; + for (int i = 0; i < 4; ++i) { + if (i < ne) { + const int e = j.grp_ptr == nullptr ? (mm + i) : j.grp_ent[mm + i]; + es[i] = e; + xr[i] = j.grp_ptr == nullptr ? 0 : (e >> 3); + } + } + float ag[4] = {0.0f, 0.0f, 0.0f, 0.0f}, au[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint kb = uint(lane); kb < nsb; kb += 32u) { + const Ds4IqSub sg = ds4_iq_sub(g, bs, int(kb), j.type); + const Ds4IqSub su = ds4_iq_sub(u, bs, int(kb), j.type); + for (int i = 0; i < 4; ++i) { + if (i < ne) { + const device int* xp = (const device int*) (j.xq + (size_t) xr[i] * (size_t) j.n_embd + + (size_t) kb * 32u); + const float dx = j.xd[(size_t) xr[i] * nsb + kb]; + ag[i] += dx * ds4_iq_dot(sg, xp); + au[i] += dx * ds4_iq_dot(su, xp); + } + } + } + for (int i = 0; i < 4; ++i) { + if (i < ne) { + const float gv = ds4_iq_lsum<1>(ag[i]), uv = ds4_iq_lsum<1>(au[i]); + if (lane == 0u) j.h[(size_t) es[i] * (size_t) j.n_ff + (size_t) row] = ds4_iq_swiglu(gv, uv, j.limit); + } + } + } +} + +// down_kernel (kDnRW = 2: a warp two rows, 16 lanes a row): part = the down against the int8 h +kernel void ds4_iq_dn(constant Ds4IqJob& j [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + ulong blob; + int m0, m1; + if (!ds4_iq_group(j, tgp.y, blob, m0, m1)) return; + const int row = (int(tgp.x) * 8 + int(tid >> 5)) * 2 + int(lane) / 16; + if (row >= j.n_embd) return; + const uint kl = uint(lane) % 16u; + const uint bs = ds4_iq_bs(j.type); + const uint nsb = uint(j.n_ff) / 32u; + const size_t rb = (size_t) (j.n_ff / 256) * (size_t) bs; + const device uchar* w = (const device uchar*) blob + (size_t) j.tr2 + (size_t) row * rb; + for (int mm = m0; mm < m1; mm += 4) { + const int ne = min(4, m1 - mm); + int es[4]; + for (int i = 0; i < 4; ++i) + if (i < ne) es[i] = j.grp_ptr == nullptr ? (mm + i) : j.grp_ent[mm + i]; + float ad[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint kb = kl; kb < nsb; kb += 16u) { + const Ds4IqSub sd = ds4_iq_sub(w, bs, int(kb), j.type); + for (int i = 0; i < 4; ++i) { + if (i < ne) { + const device int* hp = (const device int*) (j.hq + (size_t) es[i] * (size_t) j.n_ff + + (size_t) kb * 32u); + ad[i] += j.hd[(size_t) es[i] * nsb + kb] * ds4_iq_dot(sd, hp); + } + } + } + for (int i = 0; i < 4; ++i) { + if (i < ne) { + const float v = ds4_iq_lsum<2>(ad[i]); + if (kl == 0u) j.part[(size_t) es[i] * (size_t) j.n_embd + (size_t) row] = v; + } + } + } +} + +// ---------------------------------------------------------------- ds4_gguf.cu:454 mx_gate_up_kernel / :519 mx_down_kernel +// +// The MXFP4 expert GEMM (ds4_gguf.cu's moe_gate_up / moe_down / win_experts): the same frame as ds4_iq_gu / ds4_iq_dn +// above - a threadgroup of 256 is 8 warps, a warp a row of the gate and of the up (down: a row of the down), a lane +// strides the 32-value blocks of the row, and the group's entries go sixteen at a pass so that a block the weights decode +// once serves them all. A block is 17 bytes (kMxBytes): the E8M0 scale byte, then 16 bytes whose low nibble is value j and +// whose high nibble is value j + 16, each an index into the codebook {0, 1, 2, 3, 4, 6, 8, 12} and its negatives - the same +// decode as ds4_init.metal's ds4_mx_block / ds4_e8m0 (ds4_gguf.cu:143 / :154), kept as int8 codes here because the +// activations are the quantiser's int8, so the dot is integer (ds4_gguf.cu:225 mx_dot) and the two scales fold in at the +// end. +// +// The group's blob and its entries are read as ds4_iq_group above reads them (ds4_gguf.cu:429 mx_group): a plan group +// (grp_ptr null) is one expert whose entry is the group's own index and whose activation row is 0, a window group's +// entries come from grp_ent[grp_start[g] .. grp_start[g + 1]) and an entry e's activation row is e >> 3. +struct Ds4MxJob { + device const ulong* plan_ptr; // plan_ptr[g]: the expert's blob (0: the group is skipped) + device const ulong* grp_ptr; // grp_ptr[g]: a window set's blob + device const int* grp_start; // a set's first entry (n_groups + 1 of them) + device const int* grp_ent; // the entries (token * 8 + pick) + device const int* n_groups; + device const char* xq; // gate / up's int8 activations (n_embd an entry) + device const float* xd; // their scales (n_embd / 32 a row) + device const char* hq; // down's int8 input (n_ff an entry) + device const float* hd; // its scales + device float* h; // [entries][n_ff]: the clamped swiglu + device float* part; // [entries][n_embd]: each entry's down output + long tr0, tr1, tr2; // the roles' offsets in a blob + int n_embd, n_ff; + float limit; // swiglu_limit (0: none) +}; + +// the codebook of a nibble, in index order (ds4_init.metal's lut_init: {0,1,2,3,4,6,8,12}, then its negatives) +constant char ds4_mx_code[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; + +// 2^(e - 128) (ds4_init.metal:489 ds4_e8m0, ggml's GGML_E8M0_TO_FP32_HALF) +static inline float ds4_mx_e8m0(uint e) { + return as_type(e < 2u ? (0x00200000u << e) : ((e - 1u) << 23)); +} + +// a block's 16 code bytes as 8 int8 words: word j the low nibbles of bytes 4 j .. (values 4 j ..), word 4 + j their high +// nibbles (values 16 + 4 j ..) - the pairs ds4_gguf.cu:225 mx_dot reads xp[j] and xp[j + 4] with +static inline void ds4_mx_words(device const uchar* p, thread int* q) { + for (uint j = 0; j < 4u; ++j) { + uint lo = 0u, hi = 0u; + for (uint b = 0; b < 4u; ++b) { + const uint v = uint(p[1u + 4u * j + b]); + lo |= uint(uchar(ds4_mx_code[v & 15u])) << (8u * b); + hi |= uint(uchar(ds4_mx_code[v >> 4u])) << (8u * b); + } + q[j] = int(lo); + q[4u + j] = int(hi); + } +} + +// a block's 32 int8 codes against its 32 int8 activations (8 packed words, 4 values a word) +static inline int ds4_mx_dot(thread const int* q, device const int* xp) { + int s = 0; + for (uint t = 0; t < 8u; ++t) s += ds4_iq_dot4(q[t], xp[t]); + return s; +} + +// the group's blob and its entry range (ds4_iq_group's MX form, ds4_gguf.cu:429 mx_group); false: nothing to do +static inline bool ds4_mx_group(constant Ds4MxJob& j, uint g, thread ulong& blob, thread int& m0, thread int& m1) { + if (j.grp_ptr == nullptr) { + const ulong p = j.plan_ptr[g]; + if (p == 0ul) return false; + blob = p; + m0 = int(g); + m1 = m0 + 1; + return true; + } + if (int(g) >= *j.n_groups) return false; + blob = j.grp_ptr[g]; + m0 = j.grp_start[g]; + m1 = j.grp_start[g + 1]; + return true; +} + +// mx_gate_up_kernel (a warp a row of the gate and of the up, a lane a block a step): h = the clamped swiglu +kernel void ds4_mx_gu(constant Ds4MxJob& j [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + ulong blob; + int m0, m1; + if (!ds4_mx_group(j, tgp.y, blob, m0, m1)) return; + const int row = int(tgp.x) * 8 + int(tid >> 5); + if (row >= j.n_ff) return; + const uint nblk = uint(j.n_embd) / 32u; + const size_t rb = (size_t) nblk * 17u; + const device uchar* g = (const device uchar*) blob + (size_t) j.tr0 + (size_t) row * rb; + const device uchar* u = (const device uchar*) blob + (size_t) j.tr1 + (size_t) row * rb; + for (int mm = m0; mm < m1; mm += 16) { + const int ne = min(16, m1 - mm); + int es[16], xr[16]; + for (int i = 0; i < 16; ++i) { + if (i < ne) { + const int e = j.grp_ptr == nullptr ? (mm + i) : j.grp_ent[mm + i]; + es[i] = e; + xr[i] = j.grp_ptr == nullptr ? 0 : (e >> 3); + } + } + float ag[16], au[16]; + for (int i = 0; i < 16; ++i) { + ag[i] = 0.0f; + au[i] = 0.0f; + } + for (uint b = uint(lane); b < nblk; b += 32u) { + const device uchar* pg = g + (size_t) b * 17u; + const device uchar* pu = u + (size_t) b * 17u; + int wg[8], wu[8]; + ds4_mx_words(pg, wg); + ds4_mx_words(pu, wu); + const float dg = ds4_mx_e8m0(uint(pg[0])), du = ds4_mx_e8m0(uint(pu[0])); + for (int i = 0; i < 16; ++i) { + if (i < ne) { + const device int* xp = + (const device int*) (j.xq + (size_t) xr[i] * (size_t) j.n_embd + (size_t) b * 32u); + const float dx = j.xd[(size_t) xr[i] * nblk + b]; + ag[i] = fma(dg * dx, float(ds4_mx_dot(wg, xp)), ag[i]); + au[i] = fma(du * dx, float(ds4_mx_dot(wu, xp)), au[i]); + } + } + } + for (int i = 0; i < 16; ++i) { + if (i < ne) { + const float gv = ds4_iq_lsum<1>(ag[i]), uv = ds4_iq_lsum<1>(au[i]); + if (lane == 0u) j.h[(size_t) es[i] * (size_t) j.n_ff + (size_t) row] = ds4_iq_swiglu(gv, uv, j.limit); + } + } + } +} + +// mx_down_kernel (a warp a row of the down, a lane a block a step): part = the down against the int8 h +kernel void ds4_mx_dn(constant Ds4MxJob& j [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + ulong blob; + int m0, m1; + if (!ds4_mx_group(j, tgp.y, blob, m0, m1)) return; + const int row = int(tgp.x) * 8 + int(tid >> 5); + if (row >= j.n_embd) return; + const uint nblk = uint(j.n_ff) / 32u; + const size_t rb = (size_t) nblk * 17u; + const device uchar* w = (const device uchar*) blob + (size_t) j.tr2 + (size_t) row * rb; + for (int mm = m0; mm < m1; mm += 16) { + const int ne = min(16, m1 - mm); + int es[16]; + for (int i = 0; i < 16; ++i) + if (i < ne) es[i] = j.grp_ptr == nullptr ? (mm + i) : j.grp_ent[mm + i]; + float ad[16]; + for (int i = 0; i < 16; ++i) ad[i] = 0.0f; + for (uint b = uint(lane); b < nblk; b += 32u) { + const device uchar* pw = w + (size_t) b * 17u; + int ww[8]; + ds4_mx_words(pw, ww); + const float dd = ds4_mx_e8m0(uint(pw[0])); + for (int i = 0; i < 16; ++i) { + if (i < ne) { + const device int* hp = + (const device int*) (j.hq + (size_t) es[i] * (size_t) j.n_ff + (size_t) b * 32u); + ad[i] = fma(dd * j.hd[(size_t) es[i] * nblk + b], float(ds4_mx_dot(ww, hp)), ad[i]); + } + } + } + for (int i = 0; i < 16; ++i) { + if (i < ne) { + const float v = ds4_iq_lsum<1>(ad[i]); + if (lane == 0u) j.part[(size_t) es[i] * (size_t) j.n_embd + (size_t) row] = v; + } + } + } +} + +// --------------------------------------------------------------------------------------------------------------------- +// One-activation matvecs (the decode step): every lane of the row's simdgroup works, on a 32-value sub-block (IQ2_XXS) or +// a 64-value pair (Q5_K), where the 16-column kernels above give a whole 256-value block to a lane (a k = 1280 row used 5 +// lanes of 32) and read it a byte at a time. Same results up to the order of the float sums. + +kernel void ds4_q5k_mv1(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nblk = uint(b.k) / 256u; + const device uchar* w = b.w + (size_t) row * b.ld; + const uint j = lane & 3u; // this lane's pair of sub-blocks: 2j and 2j+1 (one 32-byte run of qs, both nibbles) + const uint bsel = lane >> 2; // eight blocks a pass + float acc = 0.0f; + for (uint blk = bsel; blk < nblk; blk += 8u) { + const device uchar* p = w + (size_t) blk * 176u; + const float d = float(as_type(*(const device ushort*) p)); + const float dmin = float(as_type(*(const device ushort*) (p + 2))); + int sc0, mm0, sc1, mm1; + ds4_k_scale(p + 4, int(2u * j), sc0, mm0); + ds4_k_scale(p + 4, int(2u * j + 1u), sc1, mm1); + const device uint4* qs4 = (const device uint4*) (p + 48u + 32u * j); + const device uint4* qh4 = (const device uint4*) (p + 16); + const device uint4* xq4 = (const device uint4*) (b.xq + (size_t) blk * 256u + 64u * j); + const uint4 qa = qs4[0], qb = qs4[1], ha = qh4[0], hb = qh4[1]; + const uint4 xa = xq4[0], xb = xq4[1], xc = xq4[2], xd4 = xq4[3]; // sub-block 2j: xa xb; 2j+1: xc xd4 + const uint qsw[8] = {qa.x, qa.y, qa.z, qa.w, qb.x, qb.y, qb.z, qb.w}; + const uint qhw[8] = {ha.x, ha.y, ha.z, ha.w, hb.x, hb.y, hb.z, hb.w}; + const uint xlo[8] = {xa.x, xa.y, xa.z, xa.w, xb.x, xb.y, xb.z, xb.w}; + const uint xhi[8] = {xc.x, xc.y, xc.z, xc.w, xd4.x, xd4.y, xd4.z, xd4.w}; + int sq0 = 0, sx0 = 0, sq1 = 0, sx1 = 0; + for (uint t = 0; t < 8u; ++t) { + const uint q0 = (qsw[t] & 0x0f0f0f0fu) | (((qhw[t] >> (2u * j)) & 0x01010101u) << 4); + const uint q1 = ((qsw[t] >> 4) & 0x0f0f0f0fu) | (((qhw[t] >> (2u * j + 1u)) & 0x01010101u) << 4); + for (uint i = 0; i < 4u; ++i) { + const int x0 = int(char(xlo[t] >> (8u * i))), x1 = int(char(xhi[t] >> (8u * i))); + sq0 += int((q0 >> (8u * i)) & 0xffu) * x0; + sq1 += int((q1 >> (8u * i)) & 0xffu) * x1; + sx0 += x0; + sx1 += x1; + } + } + acc = fma(b.xd[blk * 8u + 2u * j], d * float(sc0) * float(sq0) - dmin * float(mm0) * float(sx0), acc); + acc = fma(b.xd[blk * 8u + 2u * j + 1u], d * float(sc1) * float(sq1) - dmin * float(mm1) * float(sx1), acc); + } + const float v = simd_sum(acc); + if (lane == 0u) b.y[row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); +} + +kernel void ds4_iq2xxs_mv1(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nblk = uint(b.k) / 256u; + const device uchar* w = b.w + (size_t) row * b.ld; + const uint ib = lane & 7u; // this lane's 32-value sub-block + const uint bsel = lane >> 3; // four blocks a pass + float acc = 0.0f; + for (uint blk = bsel; blk < nblk; blk += 4u) { + const device uchar* p = w + (size_t) blk * 66u; + const float d = float(as_type(*(const device ushort*) p)); + const device ushort* us = (const device ushort*) (p + 2u + 8u * ib); // (2-byte aligned: four 16-bit loads) + const uint aux0 = uint(us[0]) | (uint(us[1]) << 16); + const uint aux1 = uint(us[2]) | (uint(us[3]) << 16); + const float db = d * (0.5f + float(aux1 >> 28)) * 0.25f; + const device uint4* xq4 = (const device uint4*) (b.xq + (size_t) blk * 256u + 32u * ib); + const uint4 xa = xq4[0], xb = xq4[1]; + const uint xw[8] = {xa.x, xa.y, xa.z, xa.w, xb.x, xb.y, xb.z, xb.w}; + int sum = 0; + for (uint l = 0; l < 4u; ++l) { + const ulong g = iq2xxs_grid[(aux0 >> (8u * l)) & 0xffu]; + const uint sg = uint(ksigns_iq2xs[(aux1 >> (7u * l)) & 127u]); + for (uint h = 0; h < 2u; ++h) { + const uint gw = uint(g >> (32u * h)); + const uint x = xw[2u * l + h]; + for (uint i = 0; i < 4u; ++i) { + const int code = int((gw >> (8u * i)) & 0xffu); + const int xv = int(char(x >> (8u * i))); + sum += ((sg >> (4u * h + i)) & 1u) != 0u ? -code * xv : code * xv; + } + } + } + acc = fma(b.xd[blk * 8u + ib] * db, float(sum), acc); + } + const float v = simd_sum(acc); + if (lane == 0u) b.y[row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); +} + +// The other low-bit expert types, one activation (see ds4_iq2xxs_mv1): a lane a 32-value sub-block, four blocks a pass. + +kernel void ds4_iq2xs_mv1(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nblk = uint(b.k) / 256u; + const device uchar* w = b.w + (size_t) row * b.ld; + const uint ib = lane & 7u; + const uint bsel = lane >> 3; + float acc = 0.0f; + for (uint blk = bsel; blk < nblk; blk += 4u) { + const device uchar* p = w + (size_t) blk * 74u; + const float d = float(as_type(*(const device ushort*) p)); + const device ushort* qs = (const device ushort*) (p + 2u + 8u * ib); + const uint sc = uint(p[66u + ib]); + const float db0 = d * (0.5f + float(sc & 0xfu)) * 0.25f, db1 = d * (0.5f + float(sc >> 4)) * 0.25f; + const device uint4* xq4 = (const device uint4*) (b.xq + (size_t) blk * 256u + 32u * ib); + const uint4 xa = xq4[0], xb = xq4[1]; + const uint xw[8] = {xa.x, xa.y, xa.z, xa.w, xb.x, xb.y, xb.z, xb.w}; + int s0 = 0, s1 = 0; + for (uint l = 0; l < 4u; ++l) { + const uint q = uint(qs[l]); + const ulong g = iq2xs_grid[q & 511u]; + const uint sg = uint(ksigns_iq2xs[(q >> 9) & 127u]); + int sum = 0; + for (uint h = 0; h < 2u; ++h) { + const uint gw = uint(g >> (32u * h)); + const uint x = xw[2u * l + h]; + for (uint i = 0; i < 4u; ++i) { + const int code = int((gw >> (8u * i)) & 0xffu); + const int xv = int(char(x >> (8u * i))); + sum += ((sg >> (4u * h + i)) & 1u) != 0u ? -code * xv : code * xv; + } + } + if (l < 2u) s0 += sum; else s1 += sum; + } + acc = fma(b.xd[blk * 8u + ib], db0 * float(s0) + db1 * float(s1), acc); + } + const float v = simd_sum(acc); + if (lane == 0u) b.y[row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); +} + +kernel void ds4_iq2s_mv1(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nblk = uint(b.k) / 256u; + const device uchar* w = b.w + (size_t) row * b.ld; + const uint ib = lane & 7u; + const uint bsel = lane >> 3; + float acc = 0.0f; + for (uint blk = bsel; blk < nblk; blk += 4u) { + const device uchar* p = w + (size_t) blk * 82u; + const float d = float(as_type(*(const device ushort*) p)); + const device uchar* qs = p + 2u + 4u * ib; + const device uchar* sgb = p + 34u + 4u * ib; + const uint qh = uint(p[66u + ib]); + const uint sc = uint(p[74u + ib]); + const float db0 = d * (0.5f + float(sc & 0xfu)) * 0.25f, db1 = d * (0.5f + float(sc >> 4)) * 0.25f; + const device uint4* xq4 = (const device uint4*) (b.xq + (size_t) blk * 256u + 32u * ib); + const uint4 xa = xq4[0], xb = xq4[1]; + const uint xw[8] = {xa.x, xa.y, xa.z, xa.w, xb.x, xb.y, xb.z, xb.w}; + int s0 = 0, s1 = 0; + for (uint l = 0; l < 4u; ++l) { + const uint idx = uint(qs[l]) | ((qh << (8u - 2u * l)) & 0x300u); + const ulong g = iq2s_grid[idx]; + const uint sg = uint(sgb[l]); + int sum = 0; + for (uint h = 0; h < 2u; ++h) { + const uint gw = uint(g >> (32u * h)); + const uint x = xw[2u * l + h]; + for (uint i = 0; i < 4u; ++i) { + const int code = int((gw >> (8u * i)) & 0xffu); + const int xv = int(char(x >> (8u * i))); + sum += ((sg >> (4u * h + i)) & 1u) != 0u ? -code * xv : code * xv; + } + } + if (l < 2u) s0 += sum; else s1 += sum; + } + acc = fma(b.xd[blk * 8u + ib], db0 * float(s0) + db1 * float(s1), acc); + } + const float v = simd_sum(acc); + if (lane == 0u) b.y[row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); +} + +kernel void ds4_iq3xxs_mv1(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nblk = uint(b.k) / 256u; + const device uchar* w = b.w + (size_t) row * b.ld; + const uint ib = lane & 7u; + const uint bsel = lane >> 3; + float acc = 0.0f; + for (uint blk = bsel; blk < nblk; blk += 4u) { + const device uchar* p = w + (size_t) blk * 98u; + const float d = float(as_type(*(const device ushort*) p)); + const device uchar* qs = p + 2u; + const uint w0 = uint(qs[8u * ib]) | (uint(qs[8u * ib + 1u]) << 8) | (uint(qs[8u * ib + 2u]) << 16) | (uint(qs[8u * ib + 3u]) << 24); + const uint w1 = uint(qs[8u * ib + 4u]) | (uint(qs[8u * ib + 5u]) << 8) | (uint(qs[8u * ib + 6u]) << 16) | (uint(qs[8u * ib + 7u]) << 24); + const uint aux = uint(qs[64u + 4u * ib]) | (uint(qs[65u + 4u * ib]) << 8) | (uint(qs[66u + 4u * ib]) << 16) | (uint(qs[67u + 4u * ib]) << 24); + const float db = d * (0.5f + float(aux >> 28)) * 0.5f; + const device uint4* xq4 = (const device uint4*) (b.xq + (size_t) blk * 256u + 32u * ib); + const uint4 xa = xq4[0], xb = xq4[1]; + const uint xw[8] = {xa.x, xa.y, xa.z, xa.w, xb.x, xb.y, xb.z, xb.w}; + int sum = 0; + for (uint l = 0; l < 4u; ++l) { + const uint wv = l < 2u ? w0 : w1, sh = 16u * (l & 1u); + const uint sgn = uint(ksigns_iq2xs[(aux >> (7u * l)) & 127u]); + const uint g0 = iq3xxs_grid[(wv >> sh) & 255u], g1 = iq3xxs_grid[(wv >> (sh + 8u)) & 255u]; + const uint xl = xw[2u * l], xh = xw[2u * l + 1u]; + for (uint i = 0; i < 4u; ++i) { + int v0 = int((g0 >> (8u * i)) & 0xffu), v1 = int((g1 >> (8u * i)) & 0xffu); + if ((sgn & (1u << i)) != 0u) v0 = -v0; + if ((sgn & (1u << (i + 4u))) != 0u) v1 = -v1; + sum += v0 * int(char(xl >> (8u * i))) + v1 * int(char(xh >> (8u * i))); + } + } + acc = fma(b.xd[blk * 8u + ib] * db, float(sum), acc); + } + const float v = simd_sum(acc); + if (lane == 0u) b.y[row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); +} + +// --------------------------------------------------------------------------------------------------------------------- +// Matvecs for NC activations at once (1, 2, 4, 8): the prompt windows' and the drafter's verify windows' columns. A lane takes a +// 32-value sub-block (a 64-value pair for Q5_K) of the row, decodes its weights ONCE and dots them with each of the NC columns' +// int8 activations, so the weight bytes (the cost) are read once for the window. ncol <= NC (the rest of the accumulators idle). + +template +static inline void ds4_mvn_store(const device Ds4MvJob& b, uint row, ushort lane, thread float (&acc)[NC]) { + for (uint c = 0; c < NC; ++c) { + const float v = simd_sum(acc[c]); + if (lane == 0u && int(c) < b.ncol) b.y[(size_t) c * b.ycs + row] = b.alpha * v + (b.bias != nullptr ? b.bias[row] : 0.0f); + } +} + +template +kernel void ds4_q5k_mvn(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nblk = uint(b.k) / 256u, nsb = uint(b.k) / 32u; + const device uchar* w = b.w + (size_t) row * b.ld; + const uint j = lane & 3u; + const uint bsel = lane >> 2; + float acc[NC]; + for (uint c = 0; c < NC; ++c) acc[c] = 0.0f; + for (uint blk = bsel; blk < nblk; blk += 8u) { + const device uchar* p = w + (size_t) blk * 176u; + const float d = float(as_type(*(const device ushort*) p)); + const float dmin = float(as_type(*(const device ushort*) (p + 2))); + int sc0, mm0, sc1, mm1; + ds4_k_scale(p + 4, int(2u * j), sc0, mm0); + ds4_k_scale(p + 4, int(2u * j + 1u), sc1, mm1); + const device uint4* qs4 = (const device uint4*) (p + 48u + 32u * j); + const device uint4* qh4 = (const device uint4*) (p + 16); + const uint4 qa = qs4[0], qb = qs4[1], ha = qh4[0], hb = qh4[1]; + const uint qsw[8] = {qa.x, qa.y, qa.z, qa.w, qb.x, qb.y, qb.z, qb.w}; + const uint qhw[8] = {ha.x, ha.y, ha.z, ha.w, hb.x, hb.y, hb.z, hb.w}; + uint q0w[8], q1w[8]; + for (uint t = 0; t < 8u; ++t) { + q0w[t] = (qsw[t] & 0x0f0f0f0fu) | (((qhw[t] >> (2u * j)) & 0x01010101u) << 4); + q1w[t] = ((qsw[t] >> 4) & 0x0f0f0f0fu) | (((qhw[t] >> (2u * j + 1u)) & 0x01010101u) << 4); + } + for (uint c = 0; c < NC; ++c) { + const device uint4* xq4 = (const device uint4*) (b.xq + (size_t) c * b.k + (size_t) blk * 256u + 64u * j); + const uint4 xa = xq4[0], xb = xq4[1], xc = xq4[2], xd4 = xq4[3]; + const uint xlo[8] = {xa.x, xa.y, xa.z, xa.w, xb.x, xb.y, xb.z, xb.w}; + const uint xhi[8] = {xc.x, xc.y, xc.z, xc.w, xd4.x, xd4.y, xd4.z, xd4.w}; + int sq0 = 0, sx0 = 0, sq1 = 0, sx1 = 0; + for (uint t = 0; t < 8u; ++t) { + for (uint i = 0; i < 4u; ++i) { + const int x0 = int(char(xlo[t] >> (8u * i))), x1 = int(char(xhi[t] >> (8u * i))); + sq0 += int((q0w[t] >> (8u * i)) & 0xffu) * x0; + sq1 += int((q1w[t] >> (8u * i)) & 0xffu) * x1; + sx0 += x0; + sx1 += x1; + } + } + acc[c] = fma(b.xd[(size_t) c * nsb + blk * 8u + 2u * j], d * float(sc0) * float(sq0) - dmin * float(mm0) * float(sx0), acc[c]); + acc[c] = fma(b.xd[(size_t) c * nsb + blk * 8u + 2u * j + 1u], d * float(sc1) * float(sq1) - dmin * float(mm1) * float(sx1), acc[c]); + } + } + ds4_mvn_store(b, row, lane, acc); +} + +// Q6_K (210-byte blocks; the Unsloth sets' attention, shared-expert, head and engram_kv matrices) for NC activations: four lanes a +// block, a lane one 128-value half and one of its two quarter pairs (quarters p and p + 2 share the 32 ql bytes at 32 p: the low +// nibbles are quarter p, the high ones p + 2; quarter q's 2 high bits are qh >> 2 q), eight blocks at once. The loads go through +// packed_uchar4: a block starts only 2-byte aligned. ggml's dequantize_row_q6_K arithmetic: d * scale * (q - 32). +template +kernel void ds4_q6k_mvn(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nblk = uint(b.k) / 256u, nsb = uint(b.k) / 32u; + const device uchar* w = b.w + (size_t) row * b.ld; + const uint j = lane & 3u, h = j >> 1, pr = j & 1u; + const uint bsel = lane >> 2; + float acc[NC]; + for (uint c = 0; c < NC; ++c) acc[c] = 0.0f; + for (uint blk = bsel; blk < nblk; blk += 8u) { + const device uchar* p = w + (size_t) blk * 210u; + const float d = float(as_type(*(const device ushort*) (p + 208))); + const device uchar* ql = p + 64u * h + 32u * pr; + const device uchar* qh = p + 128u + 32u * h; + const device char* scl = (const device char*) (p + 192u + 8u * h); + uint qlw[8], qhw[8]; + for (uint t = 0; t < 8u; ++t) { + qlw[t] = as_type(uchar4(*(const device packed_uchar4*) (ql + 4u * t))); + qhw[t] = as_type(uchar4(*(const device packed_uchar4*) (qh + 4u * t))); + } + for (uint qi = 0; qi < 2u; ++qi) { + const uint q = pr + 2u * qi; + const uint sh = 4u * qi; + int cw[32]; + for (uint t = 0; t < 8u; ++t) { + const uint x = ((qlw[t] >> sh) & 0x0f0f0f0fu) | (((qhw[t] >> (2u * q)) & 0x03030303u) << 4); + for (uint i = 0; i < 4u; ++i) cw[t * 4u + i] = int((x >> (8u * i)) & 0xffu) - 32; + } + const float s0 = float(scl[2u * q]), s1 = float(scl[2u * q + 1u]); + const uint g = blk * 8u + 4u * h + q; + for (uint c = 0; c < NC; ++c) { + const device uint4* xq4 = (const device uint4*) (b.xq + (size_t) c * b.k + (size_t) blk * 256u + 128u * h + 32u * q); + const uint4 xa = xq4[0], xb = xq4[1]; + const uint xw[8] = {xa.x, xa.y, xa.z, xa.w, xb.x, xb.y, xb.z, xb.w}; + int sa = 0, sb = 0; + for (uint t = 0; t < 8u; ++t) { + int part = 0; + for (uint i = 0; i < 4u; ++i) part += cw[t * 4u + i] * int(char(xw[t] >> (8u * i))); + if (t < 4u) sa += part; else sb += part; + } + acc[c] = fma(b.xd[(size_t) c * nsb + g] * d, s0 * float(sa) + s1 * float(sb), acc[c]); + } + } + } + ds4_mvn_store(b, row, lane, acc); +} + +// the grid-coded types: decode the sub-block's 32 weights (signs applied) into w[32] ... +static inline void ds4_w32_iq2xxs(const device uchar* p, uint ib, thread int (&w)[32], thread float& db) { + const float d = float(as_type(*(const device ushort*) p)); + const device ushort* us = (const device ushort*) (p + 2u + 8u * ib); + const uint aux0 = uint(us[0]) | (uint(us[1]) << 16); + const uint aux1 = uint(us[2]) | (uint(us[3]) << 16); + db = d * (0.5f + float(aux1 >> 28)) * 0.25f; + for (uint l = 0; l < 4u; ++l) { + const ulong g = iq2xxs_grid[(aux0 >> (8u * l)) & 0xffu]; + const uint sg = uint(ksigns_iq2xs[(aux1 >> (7u * l)) & 127u]); + for (uint i = 0; i < 8u; ++i) { + const int code = int((g >> (8u * i)) & 0xffu); + w[l * 8u + i] = ((sg >> i) & 1u) != 0u ? -code : code; + } + } +} +static inline void ds4_w32_iq2xs(const device uchar* p, uint ib, thread int (&w)[32], thread float& db0, thread float& db1) { + const float d = float(as_type(*(const device ushort*) p)); + const device ushort* qs = (const device ushort*) (p + 2u + 8u * ib); + const uint sc = uint(p[66u + ib]); + db0 = d * (0.5f + float(sc & 0xfu)) * 0.25f; + db1 = d * (0.5f + float(sc >> 4)) * 0.25f; + for (uint l = 0; l < 4u; ++l) { + const uint q = uint(qs[l]); + const ulong g = iq2xs_grid[q & 511u]; + const uint sg = uint(ksigns_iq2xs[(q >> 9) & 127u]); + for (uint i = 0; i < 8u; ++i) { + const int code = int((g >> (8u * i)) & 0xffu); + w[l * 8u + i] = ((sg >> i) & 1u) != 0u ? -code : code; + } + } +} +static inline void ds4_w32_iq2s(const device uchar* p, uint ib, thread int (&w)[32], thread float& db0, thread float& db1) { + const float d = float(as_type(*(const device ushort*) p)); + const device uchar* qs = p + 2u + 4u * ib; + const device uchar* sgb = p + 34u + 4u * ib; + const uint qh = uint(p[66u + ib]); + const uint sc = uint(p[74u + ib]); + db0 = d * (0.5f + float(sc & 0xfu)) * 0.25f; + db1 = d * (0.5f + float(sc >> 4)) * 0.25f; + for (uint l = 0; l < 4u; ++l) { + const uint idx = uint(qs[l]) | ((qh << (8u - 2u * l)) & 0x300u); + const ulong g = iq2s_grid[idx]; + const uint sg = uint(sgb[l]); + for (uint i = 0; i < 8u; ++i) { + const int code = int((g >> (8u * i)) & 0xffu); + w[l * 8u + i] = ((sg >> i) & 1u) != 0u ? -code : code; + } + } +} +static inline void ds4_w32_iq3xxs(const device uchar* p, uint ib, thread int (&w)[32], thread float& db) { + const float d = float(as_type(*(const device ushort*) p)); + const device uchar* qs = p + 2u; + const uint w0 = uint(qs[8u * ib]) | (uint(qs[8u * ib + 1u]) << 8) | (uint(qs[8u * ib + 2u]) << 16) | (uint(qs[8u * ib + 3u]) << 24); + const uint w1 = uint(qs[8u * ib + 4u]) | (uint(qs[8u * ib + 5u]) << 8) | (uint(qs[8u * ib + 6u]) << 16) | (uint(qs[8u * ib + 7u]) << 24); + const uint aux = uint(qs[64u + 4u * ib]) | (uint(qs[65u + 4u * ib]) << 8) | (uint(qs[66u + 4u * ib]) << 16) | (uint(qs[67u + 4u * ib]) << 24); + db = d * (0.5f + float(aux >> 28)) * 0.5f; + for (uint l = 0; l < 4u; ++l) { + const uint wv = l < 2u ? w0 : w1, sh = 16u * (l & 1u); + const uint sgn = uint(ksigns_iq2xs[(aux >> (7u * l)) & 127u]); + const uint g0 = iq3xxs_grid[(wv >> sh) & 255u], g1 = iq3xxs_grid[(wv >> (sh + 8u)) & 255u]; + for (uint i = 0; i < 4u; ++i) { + int v0 = int((g0 >> (8u * i)) & 0xffu), v1 = int((g1 >> (8u * i)) & 0xffu); + if ((sgn & (1u << i)) != 0u) v0 = -v0; + if ((sgn & (1u << (i + 4u))) != 0u) v1 = -v1; + w[l * 8u + i] = v0; + w[l * 8u + 4u + i] = v1; + } + } +} + + +// IQ1_S (UD-IQ1_S / UD-IQ1_M experts; 50-byte blocks): f16 d, qs[32] (4 grid bytes a 32-value sub-block), qh[8] u16 (3 bits a +// group of 8 more grid index, bits 12-14 the sub-block's 3-bit scale, bit 15 the sign of the +-0.125 delta). ggml's +// dequantize_row_iq1_s: dl = d (2 s + 1), w = dl (grid + delta). The integer form here (8 grid +- 1, scale / 8) is the same +// number with no rounding difference. +static inline void ds4_w32_iq1s(const device uchar* p, uint ib, thread int (&w)[32], thread float& db0, thread float& db1) { + const float d = float(as_type(*(const device ushort*) p)); + const uint qh = uint(p[34u + 2u * ib]) | (uint(p[35u + 2u * ib]) << 8); + db0 = d * float(2u * ((qh >> 12) & 7u) + 1u) * 0.125f; + db1 = db0; + const int dlt = (qh & 0x8000u) != 0u ? -1 : 1; + for (uint l = 0; l < 4u; ++l) { + const uint idx = uint(p[2u + 4u * ib + l]) | (((qh >> (3u * l)) & 7u) << 8); + const ulong g = iq1s_grid[idx]; + for (uint j = 0; j < 8u; ++j) { + const uint code = uint((g >> (8u * j)) & 0xffu); + const int c8 = (code & 0x80u) != 0u ? int(code) - 256 : int(code); // -1 / 0 / +1 + w[l * 8u + j] = 8 * c8 + dlt; + } + } +} + +// IQ1_M (the 2.00 bpw pack's 1.75-bit type): the block has no d field - its F16 scale is the top nibbles of its four +// u16 scales - and each 32-value sub-block splits into two 16-value halves with their own 3-bit scale (ggml's +// dequantize_row_iq1_m: dl1 for the first two 8-value groups, dl2 for the last two, and a +-0.125 delta per group). +// The weights are the grid's int8 codes plus that delta; the integer form here (8*code +- 1, scale / 8) is the same +// number with no rounding difference, and keeps the dot in integers like the other types. +static inline ushort ds4_iq1m_scale(const device uchar* p) { + const uint s0 = uint(p[48]) | (uint(p[49]) << 8), s1 = uint(p[50]) | (uint(p[51]) << 8); + const uint s2 = uint(p[52]) | (uint(p[53]) << 8), s3 = uint(p[54]) | (uint(p[55]) << 8); + return ushort((s0 >> 12) | ((s1 >> 8) & 0x00f0u) | ((s2 >> 4) & 0x0f00u) | (s3 & 0xf000u)); +} +static inline void ds4_w32_iq1m(const device uchar* p, uint ib, thread int (&w)[32], thread float& db0, thread float& db1) { + const float d = float(as_type(ds4_iq1m_scale(p))); + const ushort sc = ((const device ushort*) (p + 48))[ib >> 1]; + db0 = d * float(2u * ((sc >> (6u * (ib & 1u))) & 7u) + 1u) * 0.125f; + db1 = d * float(2u * ((sc >> (6u * (ib & 1u) + 3u)) & 7u) + 1u) * 0.125f; + const uint qh0 = p[32u + 2u * ib], qh1 = p[33u + 2u * ib]; + for (uint l = 0; l < 4u; ++l) { + const uint qhl = l < 2u ? qh0 : qh1; + const uint idx = uint(p[4u * ib + l]) | ((l & 1u) ? ((qhl << 4u) & 0x700u) : ((qhl << 8u) & 0x700u)); + const bool neg = (l & 1u) ? ((qhl & 0x80u) != 0u) : ((qhl & 0x08u) != 0u); + const ulong g = iq1s_grid[idx]; + for (uint j = 0; j < 8u; ++j) { + const uint code = uint((g >> (8u * j)) & 0xffu); + const int c8 = (code & 0x80u) != 0u ? int(code) - 256 : int(code); // -1 / 0 / +1 + w[l * 8u + j] = 8 * c8 + (neg ? -1 : 1); + } + } +} + +// ... and dot them with each column (the sub-block's four 8-value groups: l < 2 use db0, l >= 2 db1) +template +kernel void ds4_iqg_mvn(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nblk = uint(b.k) / 256u, nsb = uint(b.k) / 32u; + constexpr uint bs = TYPE == 16 ? 66u : (TYPE == 17 ? 74u : (TYPE == 18 ? 98u : (TYPE == 29 ? 56u : (TYPE == 19 ? 50u : 82u)))); + const device uchar* w0p = b.w + (size_t) row * b.ld; + const uint ib = lane & 7u; + const uint bsel = lane >> 3; + float acc[NC]; + for (uint c = 0; c < NC; ++c) acc[c] = 0.0f; + for (uint blk = bsel; blk < nblk; blk += 4u) { + const device uchar* p = w0p + (size_t) blk * bs; + int w[32]; + float db0, db1; + if (TYPE == 16) { ds4_w32_iq2xxs(p, ib, w, db0); db1 = db0; } + else if (TYPE == 17) ds4_w32_iq2xs(p, ib, w, db0, db1); + else if (TYPE == 22) ds4_w32_iq2s(p, ib, w, db0, db1); + else if (TYPE == 29) ds4_w32_iq1m(p, ib, w, db0, db1); + else if (TYPE == 19) ds4_w32_iq1s(p, ib, w, db0, db1); + else { ds4_w32_iq3xxs(p, ib, w, db0); db1 = db0; } + for (uint c = 0; c < NC; ++c) { + const device uint4* xq4 = (const device uint4*) (b.xq + (size_t) c * b.k + (size_t) blk * 256u + 32u * ib); + const uint4 xa = xq4[0], xb = xq4[1]; + const uint xw[8] = {xa.x, xa.y, xa.z, xa.w, xb.x, xb.y, xb.z, xb.w}; + int s0 = 0, s1 = 0; + for (uint t = 0; t < 8u; ++t) { + int part = 0; + for (uint i = 0; i < 4u; ++i) part += w[t * 4u + i] * int(char(xw[t] >> (8u * i))); + if (t < 4u) s0 += part; else s1 += part; + } + acc[c] = fma(b.xd[(size_t) c * nsb + blk * 8u + ib], db0 * float(s0) + db1 * float(s1), acc[c]); + } + } + ds4_mvn_store(b, row, lane, acc); +} + +#define DS4_MVN_INST(NC) \ + template [[host_name("ds4_q5k_mvn" #NC)]] kernel void ds4_q5k_mvn(device const uchar*, uint2, uint, ushort); \ + template [[host_name("ds4_q6k_mvn" #NC)]] kernel void ds4_q6k_mvn(device const uchar*, uint2, uint, ushort); \ + template [[host_name("ds4_iq2xxs_mvn" #NC)]] kernel void ds4_iqg_mvn(device const uchar*, uint2, uint, ushort); \ + template [[host_name("ds4_iq2xs_mvn" #NC)]] kernel void ds4_iqg_mvn(device const uchar*, uint2, uint, ushort); \ + template [[host_name("ds4_iq2s_mvn" #NC)]] kernel void ds4_iqg_mvn(device const uchar*, uint2, uint, ushort); \ + template [[host_name("ds4_iq3xxs_mvn" #NC)]] kernel void ds4_iqg_mvn(device const uchar*, uint2, uint, ushort); \ + template [[host_name("ds4_iq1m_mvn" #NC)]] kernel void ds4_iqg_mvn(device const uchar*, uint2, uint, ushort); \ + template [[host_name("ds4_iq1s_mvn" #NC)]] kernel void ds4_iqg_mvn(device const uchar*, uint2, uint, ushort); +DS4_MVN_INST(1) +DS4_MVN_INST(2) +DS4_MVN_INST(4) +DS4_MVN_INST(8) diff --git a/edge0-strata/metal/src/engine/msl/ds4_mm.metal b/edge0-strata/metal/src/engine/msl/ds4_mm.metal new file mode 100644 index 00000000..b888bdb4 --- /dev/null +++ b/edge0-strata/metal/src/engine/msl/ds4_mm.metal @@ -0,0 +1,139 @@ +// metal/src/engine/msl/ds4_mm.metal - the DeepSeek prompt chunk's expert GEMMs: moe_mm.metal's mm_gu3 / mm_down3 (matmul2d on dequantized +// weight tiles) over the engine's own expert slots and its UNPADDED sort order. Compiled after quant/units/mma/moe_mm.metal. +// tiles[i] = (slot, first row, valid rows): a tile is TM rows of one expert; rows past `valid` (the next expert's) are computed and +// dropped (the operand rows exist: the scratch carries TM rows of slack), the swiglu is DeepSeek's clamped one. + +struct DsMmArgs { + ulong gate_off; + ulong up_off; + ulong down_off; + ulong stride; // bytes between slots + uint gu_row; // bytes per gate/up row (n_embd values) + uint d_row; // bytes per down row (n_ff values) + uint n_embd; + uint n_ff; + float limit; + uint pad0; +}; + +// xs[r][c] = half(x[idx[r]][c]) for r < rows +kernel void ds4_mm_gather(device const float* x [[buffer(0)]], constant ulong& ldx [[buffer(1)]], device const int* idx [[buffer(2)]], + constant uint& n_embd [[buffer(3)]], device half* xs [[buffer(4)]], + uint2 tgp [[threadgroup_position_in_grid]], uint tid [[thread_index_in_threadgroup]]) { + const uint c = tgp.x * 256u + tid; + if (c >= n_embd) return; + xs[(ulong) tgp.y * n_embd + c] = half(x[(ulong) idx[tgp.y] * ldx + c]); +} + +template +kernel void ds4_mm_gu(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]], device half* HS [[buffer(2)]], + device const int4* tiles [[buffer(3)]], constant DsMmArgs& a [[buffer(4)]], constant uint& n_tiles [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + if (tg.y >= n_tiles) return; + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int4 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, valid = tl.z, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* E = W + ulong(e) * a.stride; + device const uchar* g0 = E + a.gate_off + ulong(n0) * a.gu_row; + device const uchar* u0 = E + a.up_off + ulong(n0) * a.gu_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tX((device half*)XS, dextents(a.n_embd, 1 << 30)); + auto rg = op.template get_right_input_cooperative_tensor(); + auto ru = op.template get_right_input_cooperative_tensor(); + auto tA = tX.template slice(0, row0); + auto accg = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + auto accu = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { accg[i] = 0.0f; accu[i] = 0.0f; } + for (int k0 = 0; k0 < int(a.n_embd); k0 += TK) { + const uint16_t cap = rg.get_capacity(), half_cap = cap / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rg.get_multidimensional_index(i); + const int k = k0 + ix[0], n = ix[1]; + device const uchar* rgw = g0 + ulong(n) * a.gu_row; + device const uchar* ruw = u0 + ulong(n) * a.gu_row; + half4 wg, wu, wg2, wu2; + W8::run(rgw, k, tab, wg, wg2); + W8::run(ruw, k, tab, wu, wu2); + rg[i] = wg.x; rg[i + 1] = wg.y; rg[i + 2] = wg.z; rg[i + 3] = wg.w; + ru[i] = wu.x; ru[i + 1] = wu.y; ru[i + 2] = wu.z; ru[i + 3] = wu.w; + const uint16_t j = i + half_cap; + rg[j] = wg2.x; rg[j + 1] = wg2.y; rg[j + 2] = wg2.z; rg[j + 3] = wg2.w; + ru[j] = wu2.x; ru[j + 1] = wu2.y; ru[j + 2] = wu2.z; ru[j + 3] = wu2.w; + } + tA = tX.template slice(k0, row0); + op.run(tA, rg, accg); + op.run(tA, ru, accu); + } + for (uint16_t i = 0; i < accg.get_capacity(); ++i) + if (accg.is_valid_element(i)) { + const float g = fmin(accg[i], a.limit), u = fmin(fmax(accu[i], -a.limit), a.limit); + accg[i] = g / (1.0f + exp(-g)) * u; + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + accg.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + if (m < valid) HS[ulong(row0 + m) * a.n_ff + n0 + f] = half(ot[sg][it]); + } +} + +template +kernel void ds4_mm_dn(device const uchar* W [[buffer(0)]], device half* YEH [[buffer(1)]], device const half* HS [[buffer(2)]], + device const int4* tiles [[buffer(3)]], constant DsMmArgs& a [[buffer(4)]], constant uint& n_tiles [[buffer(5)]], + uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) { + if (tg.y >= n_tiles || a.pad0 == 3u) return; + threadgroup float ot[SG][TM * TN]; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, SG * 32); + threadgroup_barrier(mem_flags::mem_threadgroup); + const int4 tl = tiles[tg.y]; + const int e = tl.x, row0 = tl.y, valid = tl.z, n0 = (int(tg.x) * SG + sg) * TN; + device const uchar* w0 = W + ulong(e) * a.stride + a.down_off + ulong(n0) * a.d_row; + constexpr auto desc = matmul2d_descriptor(TM, TN, TK, false, true, false, matmul2d_descriptor::mode::multiply_accumulate); + matmul2d op; + tensor, tensor_inline> tH((device half*)HS, dextents(a.n_ff, 1 << 30)); + auto rw = op.template get_right_input_cooperative_tensor(); + auto tA = tH.template slice(0, row0); + auto acc = op.template get_destination_cooperative_tensor<__tensor_ops_detail::__remove_addrspace_t, __tensor_ops_detail::__remove_addrspace_t, float>(); + for (uint16_t i = 0; i < acc.get_capacity(); ++i) + if (acc.is_valid_element(i)) acc[i] = 0.0f; + for (int k0 = 0; k0 < int(a.n_ff); k0 += TK) { + const uint16_t half_cap = rw.get_capacity() / 2; +#pragma clang loop unroll(full) + for (uint16_t i = 0; i < half_cap; i += 4) { + const auto ix = rw.get_multidimensional_index(i); + half4 w, w2; + if (a.pad0 == 1u) { w = half4(0); w2 = half4(0); } else if (a.pad0 == 2u) { w = half4(1); w2 = half4(1); W8::run(w0, 0, tab, w, w2); w = half4(0); } else + W8::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab, w, w2); + rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w; + const uint16_t j = i + half_cap; + rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w; + } + tA = tH.template slice(k0, row0); + op.run(tA, rw, acc); + } + tensor, tensor_inline> tO(ot[sg], dextents(TN, TM)); + acc.store(tO); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (int it = lane; it < TM * TN; it += 32) { + const int m = it / TN, f = it - m * TN; + if (m < valid) ((device float*) YEH)[ulong(row0 + m) * a.n_embd + n0 + f] = ot[sg][it]; + } +} + +#define DS4_MM_INST(F, T) template [[host_name("ds4_mm_gu_" #T)]] kernel decltype(ds4_mm_gu) ds4_mm_gu; template [[host_name("ds4_mm_dn_" #T)]] kernel decltype(ds4_mm_dn) ds4_mm_dn; +DS4_MM_INST(T_IQ1_M, 29) +DS4_MM_INST(T_IQ1_S, 19) +DS4_MM_INST(T_IQ2_XXS, 16) +DS4_MM_INST(T_IQ2_XS, 17) +DS4_MM_INST(T_IQ2_S, 22) +DS4_MM_INST(T_IQ3_XXS, 18) diff --git a/edge0-strata/metal/src/engine/msl/ds4_quant.metal b/edge0-strata/metal/src/engine/msl/ds4_quant.metal new file mode 100644 index 00000000..43caa393 --- /dev/null +++ b/edge0-strata/metal/src/engine/msl/ds4_quant.metal @@ -0,0 +1,32 @@ +// metal/src/engine/msl/ds4_quant.metal - the DeepSeek engine's activation quantiser, the first dsv41 kernel ported. +// +// CUDA's q8 path (ds4_gguf.cu:201, q8_quant_kernel): per 32 values d = max|x| / 127 as an f16, then q = round(x / d) clamped to +// [-128, 127] as int8. Verified bit-exact against the same arithmetic on the CPU: 4096 blocks, 131,072 values, 0 differences, +// worst scale error 4.6e-04 (half rounding) - metal/tests/dsv41/q8_quant_parity.mm. +// +// The block comes from the flat thread id (thread_position_in_grid / 32), not threadgroup_position_in_grid: the latter arrived +// as 0 in every group on this Mac's driver and made every block recompute block 0. +#include + +using namespace metal; + +// x: n values; ldx: the stride in values between consecutive blocks of 32 (32 for a contiguous row) +kernel void ds4_q8_quant(device const float* x [[buffer(0)]], constant uint& ldx [[buffer(1)]], + device char* q [[buffer(2)]], device half* d [[buffer(3)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + const uint b = gid / 32u; + const int base = int(b) * int(ldx); + const float v = x[base + int(lane)]; + float amax = fabs(v); + amax = max(amax, simd_shuffle_xor(amax, 16)); + amax = max(amax, simd_shuffle_xor(amax, 8)); + amax = max(amax, simd_shuffle_xor(amax, 4)); + amax = max(amax, simd_shuffle_xor(amax, 2)); + amax = max(amax, simd_shuffle_xor(amax, 1)); + const float scale = amax / 127.0f; + const float inv = scale > 0.0f ? 1.0f / scale : 0.0f; + int iq = int(rint(v * inv)); + iq = clamp(iq, -128, 127); + q[int(b) * 32 + int(lane)] = char(iq); + if (lane == 0) d[int(b)] = half(scale); +} diff --git a/edge0-strata/metal/src/engine/msl/ggml_grids.h b/edge0-strata/metal/src/engine/msl/ggml_grids.h new file mode 100644 index 00000000..650449e0 --- /dev/null +++ b/edge0-strata/metal/src/engine/msl/ggml_grids.h @@ -0,0 +1,21 @@ +// metal/src/engine/msl/ggml_grids.h - ggml's i-quant grid tables, for the MSL expert kernels. +// +// Generated from third_party/ggml/ggml-common.h (pinned 3cf03257f): the CUDA and Metal kernels must use bit-identical tables, +// so they are copied from the same file rather than retyped. MSL cannot include a C header, so each array becomes a +// `constant` in the address space the kernels need - same values, same order. +#pragma once + + +constant uchar kmask_iq2xs[8] = {1, 2, 4, 8, 16, 32, 64, 128}; + +constant uchar ksigns_iq2xs[128] = {0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255,}; + +constant ulong iq2xxs_grid[256] = {0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819, 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, 0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b, 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819, 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, 0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808, 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808, 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, 0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, 0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908, 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819, 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, 0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, 0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808, 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08, 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, 0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, 0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908, 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19, 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, 0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, 0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908, 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, 0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808, 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808, 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, 0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b, 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08, 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, 0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, 0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819, 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, 0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, 0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908, 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, 0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b, 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808, 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, 0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, 0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19, 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908,}; + +constant ulong iq2xs_grid[512] = {0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b192b19, 0x080808082b2b0808, 0x0808081908080819, 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x0808081908192b2b, 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, 0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, 0x08080819192b0808, 0x08080819192b2b08, 0x080808192b080819, 0x080808192b081908, 0x080808192b190808, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919, 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, 0x0808082b19080819, 0x0808082b19081908, 0x0808082b19190808, 0x0808082b19191919, 0x0808082b2b080808, 0x0808082b2b082b2b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08, 0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x0808191908080808, 0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908190819, 0x0808191908191908, 0x08081919082b0808, 0x0808191919080819, 0x0808191919081908, 0x0808191919190808, 0x08081919192b0819, 0x080819192b080808, 0x0808192b08080819, 0x0808192b08081908, 0x0808192b08190808, 0x0808192b082b192b, 0x0808192b19080808, 0x0808192b1908082b, 0x0808192b2b081908, 0x08082b0808080808, 0x08082b080808082b, 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808082b2b, 0x08082b0808190819, 0x08082b0808191908, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b0819080819, 0x08082b0819081908, 0x08082b0819190808, 0x08082b0819192b08, 0x08082b082b080808, 0x08082b082b2b0808, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, 0x08082b1908190808, 0x08082b1919080808, 0x08082b192b080819, 0x08082b192b082b19, 0x08082b2b08080808, 0x08082b2b082b0808, 0x08082b2b082b2b08, 0x08082b2b2b19192b, 0x08082b2b2b2b0808, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b, 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x0819080819080808, 0x081908081908082b, 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x08190808192b0808, 0x08190808192b2b2b, 0x081908082b080819, 0x081908082b081908, 0x081908082b190808, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919, 0x0819081908082b08, 0x0819081908190819, 0x0819081908191908, 0x08190819082b0808, 0x0819081919080819, 0x0819081919081908, 0x0819081919190808, 0x081908192b080808, 0x081908192b191908, 0x081908192b19192b, 0x0819082b08080819, 0x0819082b08081908, 0x0819082b0808192b, 0x0819082b08190808, 0x0819082b19080808, 0x0819082b192b0808, 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x08191908082b0808, 0x0819190819080819, 0x0819190819081908, 0x0819190819082b19, 0x0819190819190808, 0x08191908192b1908, 0x081919082b080808, 0x0819191908080819, 0x0819191908081908, 0x0819191908190808, 0x0819191919080808, 0x0819192b08080808, 0x0819192b08191908, 0x0819192b19082b19, 0x08192b0808080819, 0x08192b0808081908, 0x08192b0808190808, 0x08192b080819082b, 0x08192b0819080808, 0x08192b0819191908, 0x08192b082b08192b, 0x08192b1908080808, 0x08192b1908081919, 0x08192b19192b192b, 0x08192b2b19190819, 0x08192b2b2b2b2b19, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, 0x082b080808082b08, 0x082b080808082b2b, 0x082b080808190819, 0x082b080808191908, 0x082b0808082b0808, 0x082b080819080819, 0x082b080819081908, 0x082b080819190808, 0x082b08082b080808, 0x082b08082b2b0808, 0x082b081908080819, 0x082b081908081908, 0x082b081908190808, 0x082b081919080808, 0x082b081919082b08, 0x082b0819192b1919, 0x082b082b08080808, 0x082b082b082b082b, 0x082b082b2b080808, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, 0x082b190808190808, 0x082b1908082b2b19, 0x082b190819080808, 0x082b191908080808, 0x082b191919080819, 0x082b19191919082b, 0x082b19192b192b19, 0x082b192b08080819, 0x082b192b08192b2b, 0x082b192b2b2b192b, 0x082b2b0808080808, 0x082b2b0808082b08, 0x082b2b0808082b2b, 0x082b2b08082b0808, 0x082b2b0819191919, 0x082b2b082b082b08, 0x082b2b082b2b082b, 0x082b2b19192b2b08, 0x082b2b192b190808, 0x082b2b2b08082b08, 0x082b2b2b082b0808, 0x082b2b2b2b08082b, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908, 0x19080808192b0808, 0x19080808192b1919, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08, 0x1908081908190819, 0x1908081908191908, 0x19080819082b0808, 0x1908081919080819, 0x1908081919081908, 0x1908081919190808, 0x190808192b080808, 0x190808192b081919, 0x190808192b2b082b, 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b082b2b19, 0x1908082b19080808, 0x1908190808080808, 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808190819, 0x1908190808191908, 0x1908190808192b19, 0x19081908082b0808, 0x1908190819080819, 0x1908190819081908, 0x1908190819190808, 0x190819082b080808, 0x190819082b191908, 0x1908191908080819, 0x1908191908081908, 0x1908191908190808, 0x19081919082b1908, 0x1908191919080808, 0x190819192b192b2b, 0x1908192b08080808, 0x1908192b08082b2b, 0x1908192b19081908, 0x1908192b19190808, 0x19082b0808080819, 0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b0819191908, 0x19082b08192b082b, 0x19082b1908080808, 0x19082b1908190819, 0x19082b1919081908, 0x19082b1919190808, 0x19082b19192b2b19, 0x19082b2b08081908, 0x1919080808080808, 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, 0x1919080808191908, 0x19190808082b0808, 0x19190808082b2b08, 0x1919080819080819, 0x1919080819081908, 0x1919080819190808, 0x191908082b080808, 0x1919081908080819, 0x1919081908081908, 0x1919081908190808, 0x1919081908191919, 0x1919081919080808, 0x191908191908082b, 0x1919082b08080808, 0x1919082b19081908, 0x1919082b2b2b2b2b, 0x1919190808080819, 0x1919190808081908, 0x1919190808190808, 0x19191908082b0819, 0x1919190819080808, 0x19191908192b0808, 0x191919082b080819, 0x191919082b2b0819, 0x1919191908080808, 0x1919191908082b08, 0x191919192b080808, 0x191919192b082b08, 0x1919192b082b0819, 0x1919192b192b2b08, 0x1919192b2b2b0819, 0x19192b0808080808, 0x19192b0808191908, 0x19192b0819080819, 0x19192b0819190808, 0x19192b082b192b19, 0x19192b1908192b2b, 0x19192b1919080808, 0x19192b191908082b, 0x19192b2b2b081919, 0x192b080808080819, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b080819191908, 0x192b0808192b082b, 0x192b08082b08192b, 0x192b08082b2b2b19, 0x192b081908080808, 0x192b082b082b1908, 0x192b082b19082b2b, 0x192b082b2b19082b, 0x192b190808080808, 0x192b19080819192b, 0x192b191908190808, 0x192b191919080808, 0x192b191919081919, 0x192b19192b2b1908, 0x192b2b0808080819, 0x192b2b08192b2b2b, 0x192b2b19082b1919, 0x192b2b2b0808192b, 0x192b2b2b19191908, 0x192b2b2b192b082b, 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, 0x2b08080808191908, 0x2b080808082b0808, 0x2b080808082b2b2b, 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808082b080808, 0x2b0808082b08082b, 0x2b0808082b2b2b08, 0x2b0808082b2b2b2b, 0x2b08081908080819, 0x2b08081908081908, 0x2b0808190808192b, 0x2b08081908190808, 0x2b08081919080808, 0x2b08081919190819, 0x2b08081919192b19, 0x2b08082b08080808, 0x2b08082b082b0808, 0x2b08082b2b080808, 0x2b08082b2b08082b, 0x2b08082b2b2b0808, 0x2b08082b2b2b2b08, 0x2b08190808080819, 0x2b08190808081908, 0x2b08190808190808, 0x2b0819080819082b, 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, 0x2b0819082b082b19, 0x2b08191908080808, 0x2b08191919081908, 0x2b0819192b2b1919, 0x2b08192b08192b08, 0x2b08192b192b2b2b, 0x2b082b0808080808, 0x2b082b0808082b08, 0x2b082b08082b1919, 0x2b082b0819192b2b, 0x2b082b082b080808, 0x2b082b082b08082b, 0x2b082b082b2b2b08, 0x2b082b190808192b, 0x2b082b2b082b082b, 0x2b082b2b2b080808, 0x2b082b2b2b082b08, 0x2b082b2b2b19192b, 0x2b082b2b2b2b2b08, 0x2b19080808080819, 0x2b19080808081908, 0x2b19080808190808, 0x2b19080819080808, 0x2b1908081919192b, 0x2b1908082b081908, 0x2b19081908080808, 0x2b190819082b082b, 0x2b190819192b1908, 0x2b19082b1919192b, 0x2b19082b2b082b19, 0x2b19190808080808, 0x2b19190808081919, 0x2b19190819081908, 0x2b19190819190808, 0x2b19190819192b08, 0x2b191919082b2b19, 0x2b1919192b190808, 0x2b1919192b19082b, 0x2b19192b19080819, 0x2b192b0819190819, 0x2b192b082b2b192b, 0x2b192b1919082b19, 0x2b192b2b08191919, 0x2b192b2b192b0808, 0x2b2b080808080808, 0x2b2b08080808082b, 0x2b2b080808082b08, 0x2b2b080808082b2b, 0x2b2b0808082b0808, 0x2b2b0808082b2b2b, 0x2b2b08082b2b0808, 0x2b2b081919190819, 0x2b2b081919192b19, 0x2b2b08192b2b192b, 0x2b2b082b08080808, 0x2b2b082b0808082b, 0x2b2b082b08082b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b080808, 0x2b2b082b2b2b0808, 0x2b2b190819080808, 0x2b2b19082b191919, 0x2b2b192b192b1919, 0x2b2b192b2b192b08, 0x2b2b2b0808082b2b, 0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808, 0x2b2b2b082b2b2b08, 0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b, 0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b,}; + +constant ulong iq2s_grid[1024] = {0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, 0x0808080808192b19, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, 0x080808081908192b, 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x08080808192b192b, 0x08080808192b2b19, 0x080808082b080808, 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b2b0808, 0x080808082b2b1919, 0x080808082b2b2b2b, 0x0808081908080819, 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, 0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, 0x080808191919192b, 0x0808081919192b19, 0x08080819192b0808, 0x08080819192b1919, 0x08080819192b2b08, 0x080808192b080819, 0x080808192b081908, 0x080808192b190808, 0x080808192b19082b, 0x080808192b191919, 0x080808192b2b0819, 0x080808192b2b1908, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919, 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, 0x0808082b082b2b2b, 0x0808082b19080819, 0x0808082b19081908, 0x0808082b1908192b, 0x0808082b19082b19, 0x0808082b19190808, 0x0808082b19191919, 0x0808082b2b080808, 0x0808082b2b081919, 0x0808082b2b082b2b, 0x0808082b2b191908, 0x0808082b2b2b082b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, 0x0808190808082b19, 0x0808190808190808, 0x080819080819082b, 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, 0x08081908082b1908, 0x08081908082b192b, 0x08081908082b2b19, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08, 0x0808190819082b2b, 0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x0808190819192b19, 0x08081908192b0808, 0x08081908192b082b, 0x08081908192b1919, 0x080819082b080819, 0x080819082b081908, 0x080819082b08192b, 0x080819082b082b19, 0x080819082b190808, 0x080819082b191919, 0x080819082b192b08, 0x080819082b2b0819, 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908082b2b, 0x0808191908190819, 0x0808191908191908, 0x080819190819192b, 0x0808191908192b19, 0x08081919082b0808, 0x08081919082b1919, 0x08081919082b2b08, 0x0808191919080819, 0x0808191919081908, 0x080819191908192b, 0x0808191919082b19, 0x0808191919190808, 0x080819191919082b, 0x0808191919191919, 0x0808191919192b08, 0x08081919192b0819, 0x08081919192b1908, 0x080819192b080808, 0x080819192b08082b, 0x080819192b081919, 0x080819192b082b08, 0x080819192b190819, 0x080819192b191908, 0x080819192b2b0808, 0x0808192b08080819, 0x0808192b08081908, 0x0808192b0808192b, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b08191919, 0x0808192b19080808, 0x0808192b19081919, 0x0808192b19082b08, 0x0808192b19190819, 0x0808192b19191908, 0x0808192b192b0808, 0x0808192b2b080819, 0x0808192b2b081908, 0x0808192b2b190808, 0x08082b0808080808, 0x08082b080808082b, 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808190819, 0x08082b0808191908, 0x08082b080819192b, 0x08082b0808192b19, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b08082b2b2b, 0x08082b0819080819, 0x08082b0819081908, 0x08082b081908192b, 0x08082b0819082b19, 0x08082b0819190808, 0x08082b081919082b, 0x08082b0819191919, 0x08082b0819192b08, 0x08082b08192b0819, 0x08082b08192b1908, 0x08082b082b080808, 0x08082b082b081919, 0x08082b082b191908, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, 0x08082b1908190808, 0x08082b190819082b, 0x08082b1908191919, 0x08082b1908192b08, 0x08082b19082b0819, 0x08082b1919080808, 0x08082b1919081919, 0x08082b1919082b08, 0x08082b1919190819, 0x08082b1919191908, 0x08082b19192b0808, 0x08082b192b080819, 0x08082b192b190808, 0x08082b2b08080808, 0x08082b2b08190819, 0x08082b2b08191908, 0x08082b2b082b082b, 0x08082b2b082b2b08, 0x08082b2b082b2b2b, 0x08082b2b19190808, 0x08082b2b2b192b19, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b, 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x08190808082b192b, 0x0819080819080808, 0x081908081908082b, 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x081908081919192b, 0x0819080819192b19, 0x08190808192b0808, 0x08190808192b082b, 0x08190808192b1919, 0x08190808192b2b08, 0x081908082b080819, 0x081908082b081908, 0x081908082b08192b, 0x081908082b190808, 0x081908082b191919, 0x081908082b192b08, 0x081908082b2b0819, 0x081908082b2b1908, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919, 0x0819081908082b08, 0x0819081908082b2b, 0x0819081908190819, 0x0819081908191908, 0x081908190819192b, 0x0819081908192b19, 0x08190819082b0808, 0x08190819082b082b, 0x08190819082b1919, 0x08190819082b2b08, 0x0819081919080819, 0x0819081919081908, 0x081908191908192b, 0x0819081919082b19, 0x0819081919190808, 0x081908191919082b, 0x0819081919191919, 0x0819081919192b08, 0x08190819192b0819, 0x08190819192b1908, 0x081908192b080808, 0x081908192b08082b, 0x081908192b081919, 0x081908192b082b08, 0x081908192b190819, 0x081908192b191908, 0x0819082b08080819, 0x0819082b08081908, 0x0819082b08082b19, 0x0819082b08190808, 0x0819082b08191919, 0x0819082b082b0819, 0x0819082b082b1908, 0x0819082b19080808, 0x0819082b19081919, 0x0819082b19190819, 0x0819082b19191908, 0x0819082b2b080819, 0x0819082b2b081908, 0x0819082b2b190808, 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x081919080819192b, 0x0819190808192b19, 0x08191908082b0808, 0x08191908082b1919, 0x08191908082b2b08, 0x0819190819080819, 0x0819190819081908, 0x081919081908192b, 0x0819190819082b19, 0x0819190819190808, 0x081919081919082b, 0x0819190819191919, 0x0819190819192b08, 0x08191908192b0819, 0x08191908192b1908, 0x081919082b080808, 0x081919082b08082b, 0x081919082b081919, 0x081919082b082b08, 0x081919082b190819, 0x081919082b191908, 0x081919082b2b0808, 0x0819191908080819, 0x0819191908081908, 0x081919190808192b, 0x0819191908082b19, 0x0819191908190808, 0x081919190819082b, 0x0819191908191919, 0x0819191908192b08, 0x08191919082b0819, 0x08191919082b1908, 0x0819191919080808, 0x081919191908082b, 0x0819191919081919, 0x0819191919082b08, 0x0819191919190819, 0x0819191919191908, 0x08191919192b0808, 0x081919192b080819, 0x081919192b081908, 0x081919192b190808, 0x0819192b08080808, 0x0819192b08081919, 0x0819192b08082b08, 0x0819192b08190819, 0x0819192b08191908, 0x0819192b082b0808, 0x0819192b19080819, 0x0819192b19081908, 0x0819192b19190808, 0x0819192b2b080808, 0x0819192b2b2b2b2b, 0x08192b0808080819, 0x08192b0808081908, 0x08192b080808192b, 0x08192b0808082b19, 0x08192b0808190808, 0x08192b0808191919, 0x08192b0808192b08, 0x08192b08082b0819, 0x08192b0819080808, 0x08192b081908082b, 0x08192b0819081919, 0x08192b0819082b08, 0x08192b0819190819, 0x08192b0819191908, 0x08192b08192b0808, 0x08192b082b080819, 0x08192b082b081908, 0x08192b1908080808, 0x08192b190808082b, 0x08192b1908081919, 0x08192b1908082b08, 0x08192b1908190819, 0x08192b1908191908, 0x08192b19082b0808, 0x08192b1919080819, 0x08192b1919081908, 0x08192b1919190808, 0x08192b19192b2b19, 0x08192b192b2b082b, 0x08192b2b08081908, 0x08192b2b08190808, 0x08192b2b19080808, 0x08192b2b1919192b, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, 0x082b080808082b08, 0x082b080808190819, 0x082b080808191908, 0x082b08080819192b, 0x082b080808192b19, 0x082b0808082b0808, 0x082b0808082b1919, 0x082b0808082b2b2b, 0x082b080819080819, 0x082b080819081908, 0x082b080819190808, 0x082b08081919082b, 0x082b080819191919, 0x082b0808192b1908, 0x082b08082b080808, 0x082b08082b082b2b, 0x082b08082b191908, 0x082b08082b2b2b2b, 0x082b081908080819, 0x082b081908081908, 0x082b081908190808, 0x082b08190819082b, 0x082b081908191919, 0x082b0819082b0819, 0x082b081919080808, 0x082b08191908082b, 0x082b081919081919, 0x082b081919190819, 0x082b081919191908, 0x082b0819192b0808, 0x082b08192b080819, 0x082b08192b081908, 0x082b08192b190808, 0x082b082b08080808, 0x082b082b08082b2b, 0x082b082b082b082b, 0x082b082b082b2b08, 0x082b082b082b2b2b, 0x082b082b19081908, 0x082b082b19190808, 0x082b082b2b082b08, 0x082b082b2b082b2b, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, 0x082b19080808192b, 0x082b190808082b19, 0x082b190808190808, 0x082b190808191919, 0x082b190808192b08, 0x082b1908082b0819, 0x082b1908082b1908, 0x082b190819080808, 0x082b19081908082b, 0x082b190819081919, 0x082b190819082b08, 0x082b190819190819, 0x082b190819191908, 0x082b1908192b0808, 0x082b19082b080819, 0x082b19082b081908, 0x082b19082b190808, 0x082b191908080808, 0x082b191908081919, 0x082b191908082b08, 0x082b191908190819, 0x082b191908191908, 0x082b1919082b0808, 0x082b191919080819, 0x082b191919081908, 0x082b191919190808, 0x082b1919192b192b, 0x082b19192b080808, 0x082b192b08080819, 0x082b192b08081908, 0x082b192b08190808, 0x082b192b19080808, 0x082b192b19192b19, 0x082b2b0808080808, 0x082b2b0808081919, 0x082b2b0808190819, 0x082b2b0808191908, 0x082b2b0819080819, 0x082b2b0819081908, 0x082b2b0819190808, 0x082b2b082b082b2b, 0x082b2b082b2b2b2b, 0x082b2b1908080819, 0x082b2b1908081908, 0x082b2b1908190808, 0x082b2b192b191919, 0x082b2b2b08082b2b, 0x082b2b2b082b082b, 0x082b2b2b192b1908, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x1908080808192b2b, 0x19080808082b0819, 0x19080808082b1908, 0x19080808082b192b, 0x1908080819080808, 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908, 0x190808081919192b, 0x1908080819192b19, 0x19080808192b0808, 0x19080808192b082b, 0x19080808192b1919, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, 0x190808082b191919, 0x190808082b192b08, 0x190808082b2b0819, 0x190808082b2b1908, 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08, 0x1908081908190819, 0x1908081908191908, 0x190808190819192b, 0x1908081908192b19, 0x19080819082b0808, 0x19080819082b082b, 0x19080819082b1919, 0x1908081919080819, 0x1908081919081908, 0x190808191908192b, 0x1908081919082b19, 0x1908081919190808, 0x190808191919082b, 0x1908081919191919, 0x1908081919192b08, 0x19080819192b0819, 0x19080819192b1908, 0x190808192b080808, 0x190808192b08082b, 0x190808192b081919, 0x190808192b082b08, 0x190808192b190819, 0x190808192b191908, 0x190808192b2b0808, 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b08191919, 0x1908082b08192b08, 0x1908082b082b1908, 0x1908082b19080808, 0x1908082b19081919, 0x1908082b19082b08, 0x1908082b19190819, 0x1908082b19191908, 0x1908082b192b0808, 0x1908082b2b080819, 0x1908082b2b081908, 0x1908190808080808, 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808082b2b, 0x1908190808190819, 0x1908190808191908, 0x190819080819192b, 0x1908190808192b19, 0x19081908082b0808, 0x19081908082b082b, 0x19081908082b1919, 0x19081908082b2b08, 0x1908190819080819, 0x1908190819081908, 0x190819081908192b, 0x1908190819082b19, 0x1908190819190808, 0x190819081919082b, 0x1908190819191919, 0x1908190819192b08, 0x19081908192b0819, 0x19081908192b1908, 0x190819082b080808, 0x190819082b08082b, 0x190819082b081919, 0x190819082b082b08, 0x190819082b190819, 0x190819082b191908, 0x190819082b2b0808, 0x1908191908080819, 0x1908191908081908, 0x190819190808192b, 0x1908191908082b19, 0x1908191908190808, 0x190819190819082b, 0x1908191908191919, 0x1908191908192b08, 0x19081919082b0819, 0x19081919082b1908, 0x1908191919080808, 0x190819191908082b, 0x1908191919081919, 0x1908191919082b08, 0x1908191919190819, 0x1908191919191908, 0x19081919192b0808, 0x19081919192b2b2b, 0x190819192b080819, 0x190819192b081908, 0x190819192b190808, 0x1908192b08080808, 0x1908192b0808082b, 0x1908192b08081919, 0x1908192b08082b08, 0x1908192b08190819, 0x1908192b08191908, 0x1908192b082b0808, 0x1908192b19080819, 0x1908192b19081908, 0x1908192b19190808, 0x1908192b2b080808, 0x1908192b2b2b1919, 0x19082b0808080819, 0x19082b0808081908, 0x19082b0808082b19, 0x19082b0808190808, 0x19082b080819082b, 0x19082b0808191919, 0x19082b0808192b08, 0x19082b08082b0819, 0x19082b08082b1908, 0x19082b0819080808, 0x19082b081908082b, 0x19082b0819081919, 0x19082b0819082b08, 0x19082b0819190819, 0x19082b0819191908, 0x19082b08192b0808, 0x19082b082b081908, 0x19082b082b190808, 0x19082b1908080808, 0x19082b190808082b, 0x19082b1908081919, 0x19082b1908082b08, 0x19082b1908190819, 0x19082b1908191908, 0x19082b19082b0808, 0x19082b1919080819, 0x19082b1919081908, 0x19082b1919190808, 0x19082b192b080808, 0x19082b192b19192b, 0x19082b2b08080819, 0x19082b2b08081908, 0x19082b2b08190808, 0x19082b2b19080808, 0x1919080808080808, 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, 0x1919080808190819, 0x1919080808191908, 0x191908080819192b, 0x1919080808192b19, 0x19190808082b0808, 0x19190808082b082b, 0x19190808082b1919, 0x19190808082b2b08, 0x1919080819080819, 0x1919080819081908, 0x191908081908192b, 0x1919080819082b19, 0x1919080819190808, 0x191908081919082b, 0x1919080819191919, 0x1919080819192b08, 0x19190808192b0819, 0x19190808192b1908, 0x191908082b080808, 0x191908082b08082b, 0x191908082b081919, 0x191908082b082b08, 0x191908082b190819, 0x191908082b191908, 0x1919081908080819, 0x1919081908081908, 0x191908190808192b, 0x1919081908082b19, 0x1919081908190808, 0x191908190819082b, 0x1919081908191919, 0x1919081908192b08, 0x19190819082b0819, 0x19190819082b1908, 0x1919081919080808, 0x191908191908082b, 0x1919081919081919, 0x1919081919082b08, 0x1919081919190819, 0x1919081919191908, 0x19190819192b0808, 0x191908192b080819, 0x191908192b081908, 0x191908192b190808, 0x1919082b08080808, 0x1919082b08081919, 0x1919082b08082b08, 0x1919082b08190819, 0x1919082b08191908, 0x1919082b082b0808, 0x1919082b19080819, 0x1919082b19081908, 0x1919082b19190808, 0x1919082b192b2b19, 0x1919082b2b080808, 0x1919190808080819, 0x1919190808081908, 0x191919080808192b, 0x1919190808082b19, 0x1919190808190808, 0x191919080819082b, 0x1919190808191919, 0x1919190808192b08, 0x19191908082b0819, 0x19191908082b1908, 0x1919190819080808, 0x191919081908082b, 0x1919190819081919, 0x1919190819082b08, 0x1919190819190819, 0x1919190819191908, 0x19191908192b0808, 0x191919082b080819, 0x191919082b081908, 0x191919082b190808, 0x1919191908080808, 0x191919190808082b, 0x1919191908081919, 0x1919191908082b08, 0x1919191908190819, 0x1919191908191908, 0x19191919082b0808, 0x1919191919080819, 0x1919191919081908, 0x1919191919190808, 0x191919192b080808, 0x1919192b08080819, 0x1919192b08081908, 0x1919192b08190808, 0x1919192b082b192b, 0x1919192b19080808, 0x19192b0808080808, 0x19192b080808082b, 0x19192b0808081919, 0x19192b0808082b08, 0x19192b0808190819, 0x19192b0808191908, 0x19192b08082b0808, 0x19192b0819080819, 0x19192b0819081908, 0x19192b0819190808, 0x19192b0819192b2b, 0x19192b082b080808, 0x19192b1908080819, 0x19192b1908081908, 0x19192b1908190808, 0x19192b1919080808, 0x19192b2b08080808, 0x19192b2b08192b19, 0x19192b2b2b081919, 0x19192b2b2b2b2b08, 0x192b080808080819, 0x192b080808081908, 0x192b08080808192b, 0x192b080808190808, 0x192b08080819082b, 0x192b080808191919, 0x192b080808192b08, 0x192b0808082b0819, 0x192b0808082b1908, 0x192b080819080808, 0x192b080819081919, 0x192b080819082b08, 0x192b080819190819, 0x192b080819191908, 0x192b0808192b0808, 0x192b08082b081908, 0x192b08082b190808, 0x192b081908080808, 0x192b08190808082b, 0x192b081908081919, 0x192b081908082b08, 0x192b081908190819, 0x192b081908191908, 0x192b0819082b0808, 0x192b081919080819, 0x192b081919081908, 0x192b081919190808, 0x192b08192b080808, 0x192b08192b192b19, 0x192b082b08081908, 0x192b082b08190808, 0x192b082b19080808, 0x192b082b1919192b, 0x192b082b2b2b0819, 0x192b190808080808, 0x192b190808081919, 0x192b190808082b08, 0x192b190808190819, 0x192b190808191908, 0x192b1908082b0808, 0x192b190819080819, 0x192b190819081908, 0x192b190819190808, 0x192b19082b080808, 0x192b191908080819, 0x192b191908081908, 0x192b191908190808, 0x192b191919080808, 0x192b191919082b2b, 0x192b1919192b2b08, 0x192b19192b19082b, 0x192b192b08080808, 0x192b192b2b191908, 0x192b2b0808080819, 0x192b2b0808081908, 0x192b2b0808190808, 0x192b2b08192b1919, 0x192b2b082b192b08, 0x192b2b1908080808, 0x192b2b19082b2b2b, 0x192b2b2b1908082b, 0x192b2b2b2b2b0819, 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, 0x2b08080808191908, 0x2b08080808192b19, 0x2b080808082b0808, 0x2b080808082b1919, 0x2b08080819080819, 0x2b08080819081908, 0x2b08080819190808, 0x2b0808081919082b, 0x2b08080819191919, 0x2b08080819192b08, 0x2b080808192b0819, 0x2b0808082b080808, 0x2b0808082b081919, 0x2b0808082b190819, 0x2b0808082b191908, 0x2b08081908080819, 0x2b08081908081908, 0x2b08081908082b19, 0x2b08081908190808, 0x2b0808190819082b, 0x2b08081908191919, 0x2b08081908192b08, 0x2b080819082b0819, 0x2b080819082b1908, 0x2b08081919080808, 0x2b0808191908082b, 0x2b08081919081919, 0x2b08081919082b08, 0x2b08081919190819, 0x2b08081919191908, 0x2b0808192b080819, 0x2b0808192b081908, 0x2b0808192b190808, 0x2b0808192b2b2b19, 0x2b08082b08080808, 0x2b08082b08081919, 0x2b08082b08082b2b, 0x2b08082b08190819, 0x2b08082b08191908, 0x2b08082b19080819, 0x2b08082b19081908, 0x2b08082b19190808, 0x2b08190808080819, 0x2b08190808081908, 0x2b0819080808192b, 0x2b08190808082b19, 0x2b08190808190808, 0x2b0819080819082b, 0x2b08190808191919, 0x2b08190808192b08, 0x2b081908082b0819, 0x2b08190819080808, 0x2b0819081908082b, 0x2b08190819081919, 0x2b08190819082b08, 0x2b08190819190819, 0x2b08190819191908, 0x2b081908192b0808, 0x2b0819082b080819, 0x2b0819082b081908, 0x2b0819082b190808, 0x2b08191908080808, 0x2b0819190808082b, 0x2b08191908081919, 0x2b08191908082b08, 0x2b08191908190819, 0x2b08191908191908, 0x2b081919082b0808, 0x2b08191919080819, 0x2b08191919081908, 0x2b08191919190808, 0x2b0819192b080808, 0x2b0819192b082b2b, 0x2b08192b08080819, 0x2b08192b08081908, 0x2b08192b08190808, 0x2b08192b082b2b19, 0x2b08192b19080808, 0x2b082b0808080808, 0x2b082b0808081919, 0x2b082b0808190819, 0x2b082b0808191908, 0x2b082b0819080819, 0x2b082b0819081908, 0x2b082b0819190808, 0x2b082b082b2b082b, 0x2b082b1908080819, 0x2b082b1908081908, 0x2b082b1919080808, 0x2b082b19192b1919, 0x2b082b2b082b082b, 0x2b082b2b19192b08, 0x2b082b2b19192b2b, 0x2b082b2b2b08082b, 0x2b082b2b2b2b082b, 0x2b19080808080819, 0x2b19080808081908, 0x2b19080808082b19, 0x2b19080808190808, 0x2b1908080819082b, 0x2b19080808191919, 0x2b19080808192b08, 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908081908082b, 0x2b19080819081919, 0x2b19080819082b08, 0x2b19080819190819, 0x2b19080819191908, 0x2b190808192b0808, 0x2b1908082b080819, 0x2b1908082b081908, 0x2b1908082b190808, 0x2b19081908080808, 0x2b19081908081919, 0x2b19081908190819, 0x2b19081908191908, 0x2b19081919080819, 0x2b19081919081908, 0x2b19081919190808, 0x2b19081919192b2b, 0x2b19082b08080819, 0x2b19082b08081908, 0x2b19082b08190808, 0x2b19082b19080808, 0x2b19082b2b2b192b, 0x2b19190808080808, 0x2b1919080808082b, 0x2b19190808081919, 0x2b19190808082b08, 0x2b19190808190819, 0x2b19190808191908, 0x2b191908082b0808, 0x2b19190819080819, 0x2b19190819081908, 0x2b19190819190808, 0x2b1919082b080808, 0x2b1919082b19192b, 0x2b19191908080819, 0x2b19191908081908, 0x2b19191908190808, 0x2b19191919080808, 0x2b1919192b192b08, 0x2b1919192b2b0819, 0x2b19192b08080808, 0x2b19192b1908192b, 0x2b19192b192b1908, 0x2b192b0808080819, 0x2b192b0808081908, 0x2b192b0808190808, 0x2b192b08082b192b, 0x2b192b0819080808, 0x2b192b082b2b2b19, 0x2b192b1908080808, 0x2b192b1919082b19, 0x2b192b191919082b, 0x2b192b2b2b190808, 0x2b2b080808080808, 0x2b2b080808081919, 0x2b2b080808082b2b, 0x2b2b080808191908, 0x2b2b0808082b082b, 0x2b2b0808082b2b2b, 0x2b2b080819080819, 0x2b2b080819081908, 0x2b2b080819190808, 0x2b2b08082b2b082b, 0x2b2b08082b2b2b2b, 0x2b2b081919080808, 0x2b2b0819192b1919, 0x2b2b082b0808082b, 0x2b2b082b08082b2b, 0x2b2b082b082b082b, 0x2b2b082b082b2b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b08082b, 0x2b2b082b2b082b08, 0x2b2b082b2b082b2b, 0x2b2b082b2b2b2b08, 0x2b2b190808080819, 0x2b2b190808081908, 0x2b2b190808190808, 0x2b2b190819080808, 0x2b2b19082b082b19, 0x2b2b19082b2b1908, 0x2b2b191908080808, 0x2b2b191908192b19, 0x2b2b192b19190819, 0x2b2b2b0808082b2b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b082b, 0x2b2b2b1919191908, 0x2b2b2b192b08192b, 0x2b2b2b2b08082b08, 0x2b2b2b2b08082b2b, 0x2b2b2b2b082b0808, 0x2b2b2b2b082b082b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b2b082b08, 0x2b2b2b2b2b2b2b2b,}; + +constant uint iq3xxs_grid[256] = {0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414, 0x04041c0c, 0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14, 0x040c140c, 0x040c142c, 0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404, 0x04140414, 0x04140424, 0x04140c0c, 0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e, 0x04142c0c, 0x04142c3e, 0x04143e2c, 0x041c040c, 0x041c043e, 0x041c0c04, 0x041c0c14, 0x041c142c, 0x041c3e04, 0x04240c1c, 0x04241c3e, 0x04242424, 0x04242c3e, 0x04243e1c, 0x04243e2c, 0x042c040c, 0x042c043e, 0x042c1c14, 0x042c2c14, 0x04341c2c, 0x04343424, 0x043e0c04, 0x043e0c24, 0x043e0c34, 0x043e241c, 0x043e340c, 0x0c04040c, 0x0c04041c, 0x0c040c04, 0x0c040c14, 0x0c04140c, 0x0c04141c, 0x0c041c04, 0x0c041c14, 0x0c041c24, 0x0c04243e, 0x0c042c04, 0x0c0c0404, 0x0c0c0414, 0x0c0c0c0c, 0x0c0c1404, 0x0c0c1414, 0x0c14040c, 0x0c14041c, 0x0c140c04, 0x0c140c14, 0x0c14140c, 0x0c141c04, 0x0c143e14, 0x0c1c0404, 0x0c1c0414, 0x0c1c1404, 0x0c1c1c0c, 0x0c1c2434, 0x0c1c3434, 0x0c24040c, 0x0c24042c, 0x0c242c04, 0x0c2c1404, 0x0c2c1424, 0x0c2c2434, 0x0c2c3e0c, 0x0c34042c, 0x0c3e1414, 0x0c3e2404, 0x14040404, 0x14040414, 0x14040c0c, 0x14040c1c, 0x14041404, 0x14041414, 0x14041434, 0x14041c0c, 0x14042414, 0x140c040c, 0x140c041c, 0x140c042c, 0x140c0c04, 0x140c0c14, 0x140c140c, 0x140c1c04, 0x140c341c, 0x140c343e, 0x140c3e04, 0x14140404, 0x14140414, 0x14140c0c, 0x14140c3e, 0x14141404, 0x14141414, 0x14141c3e, 0x14142404, 0x14142c2c, 0x141c040c, 0x141c0c04, 0x141c0c24, 0x141c3e04, 0x141c3e24, 0x14241c2c, 0x14242c1c, 0x142c041c, 0x142c143e, 0x142c240c, 0x142c3e24, 0x143e040c, 0x143e041c, 0x143e0c34, 0x143e242c, 0x1c04040c, 0x1c040c04, 0x1c040c14, 0x1c04140c, 0x1c04141c, 0x1c042c04, 0x1c04342c, 0x1c043e14, 0x1c0c0404, 0x1c0c0414, 0x1c0c1404, 0x1c0c1c0c, 0x1c0c2424, 0x1c0c2434, 0x1c14040c, 0x1c14041c, 0x1c140c04, 0x1c14142c, 0x1c142c14, 0x1c143e14, 0x1c1c0c0c, 0x1c1c1c1c, 0x1c241c04, 0x1c24243e, 0x1c243e14, 0x1c2c0404, 0x1c2c0434, 0x1c2c1414, 0x1c2c2c2c, 0x1c340c24, 0x1c341c34, 0x1c34341c, 0x1c3e1c1c, 0x1c3e3404, 0x24040424, 0x24040c3e, 0x24041c2c, 0x24041c3e, 0x24042c1c, 0x24042c3e, 0x240c3e24, 0x24141404, 0x24141c3e, 0x24142404, 0x24143404, 0x24143434, 0x241c043e, 0x241c242c, 0x24240424, 0x24242c0c, 0x24243424, 0x242c142c, 0x242c241c, 0x242c3e04, 0x243e042c, 0x243e0c04, 0x243e0c14, 0x243e1c04, 0x2c040c14, 0x2c04240c, 0x2c043e04, 0x2c0c0404, 0x2c0c0434, 0x2c0c1434, 0x2c0c2c2c, 0x2c140c24, 0x2c141c14, 0x2c143e14, 0x2c1c0414, 0x2c1c2c1c, 0x2c240c04, 0x2c24141c, 0x2c24143e, 0x2c243e14, 0x2c2c0414, 0x2c2c1c0c, 0x2c342c04, 0x2c3e1424, 0x2c3e2414, 0x34041424, 0x34042424, 0x34042434, 0x34043424, 0x340c140c, 0x340c340c, 0x34140c3e, 0x34143424, 0x341c1c04, 0x341c1c34, 0x34242424, 0x342c042c, 0x342c2c14, 0x34341c1c, 0x343e041c, 0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14, 0x3e042c14, 0x3e0c1434, 0x3e0c2404, 0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c, 0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04,}; + +constant uint iq3s_grid[512] = {0x01010101, 0x01010103, 0x01010105, 0x0101010b, 0x0101010f, 0x01010301, 0x01010303, 0x01010305, 0x01010309, 0x0101030d, 0x01010501, 0x01010503, 0x0101050b, 0x01010707, 0x01010901, 0x01010905, 0x0101090b, 0x0101090f, 0x01010b03, 0x01010b07, 0x01010d01, 0x01010d05, 0x01010f03, 0x01010f09, 0x01010f0f, 0x01030101, 0x01030103, 0x01030105, 0x01030109, 0x01030301, 0x01030303, 0x0103030b, 0x01030501, 0x01030507, 0x0103050f, 0x01030703, 0x0103070b, 0x01030909, 0x01030d03, 0x01030d0b, 0x01030f05, 0x01050101, 0x01050103, 0x0105010b, 0x0105010f, 0x01050301, 0x01050307, 0x0105030d, 0x01050503, 0x0105050b, 0x01050701, 0x01050709, 0x01050905, 0x0105090b, 0x0105090f, 0x01050b03, 0x01050b07, 0x01050f01, 0x01050f07, 0x01070107, 0x01070303, 0x0107030b, 0x01070501, 0x01070505, 0x01070703, 0x01070707, 0x0107070d, 0x01070909, 0x01070b01, 0x01070b05, 0x01070d0f, 0x01070f03, 0x01070f0b, 0x01090101, 0x01090307, 0x0109030f, 0x01090503, 0x01090509, 0x01090705, 0x01090901, 0x01090907, 0x01090b03, 0x01090f01, 0x010b0105, 0x010b0109, 0x010b0501, 0x010b0505, 0x010b050d, 0x010b0707, 0x010b0903, 0x010b090b, 0x010b090f, 0x010b0d0d, 0x010b0f07, 0x010d010d, 0x010d0303, 0x010d0307, 0x010d0703, 0x010d0b05, 0x010d0f03, 0x010f0101, 0x010f0105, 0x010f0109, 0x010f0501, 0x010f0505, 0x010f050d, 0x010f0707, 0x010f0b01, 0x010f0b09, 0x03010101, 0x03010103, 0x03010105, 0x03010109, 0x03010301, 0x03010303, 0x03010307, 0x0301030b, 0x0301030f, 0x03010501, 0x03010505, 0x03010703, 0x03010709, 0x0301070d, 0x03010b09, 0x03010b0d, 0x03010d03, 0x03010f05, 0x03030101, 0x03030103, 0x03030107, 0x0303010d, 0x03030301, 0x03030309, 0x03030503, 0x03030701, 0x03030707, 0x03030903, 0x03030b01, 0x03030b05, 0x03030f01, 0x03030f0d, 0x03050101, 0x03050305, 0x0305030b, 0x0305030f, 0x03050501, 0x03050509, 0x03050705, 0x03050901, 0x03050907, 0x03050b0b, 0x03050d01, 0x03050f05, 0x03070103, 0x03070109, 0x0307010f, 0x03070301, 0x03070307, 0x03070503, 0x0307050f, 0x03070701, 0x03070709, 0x03070903, 0x03070d05, 0x03070f01, 0x03090107, 0x0309010b, 0x03090305, 0x03090309, 0x03090703, 0x03090707, 0x03090905, 0x0309090d, 0x03090b01, 0x03090b09, 0x030b0103, 0x030b0301, 0x030b0307, 0x030b0503, 0x030b0701, 0x030b0705, 0x030b0b03, 0x030d0501, 0x030d0509, 0x030d050f, 0x030d0909, 0x030d090d, 0x030f0103, 0x030f0107, 0x030f0301, 0x030f0305, 0x030f0503, 0x030f070b, 0x030f0903, 0x030f0d05, 0x030f0f01, 0x05010101, 0x05010103, 0x05010107, 0x0501010b, 0x0501010f, 0x05010301, 0x05010305, 0x05010309, 0x0501030d, 0x05010503, 0x05010507, 0x0501050f, 0x05010701, 0x05010705, 0x05010903, 0x05010907, 0x0501090b, 0x05010b01, 0x05010b05, 0x05010d0f, 0x05010f01, 0x05010f07, 0x05010f0b, 0x05030101, 0x05030105, 0x05030301, 0x05030307, 0x0503030f, 0x05030505, 0x0503050b, 0x05030703, 0x05030709, 0x05030905, 0x05030b03, 0x05050103, 0x05050109, 0x0505010f, 0x05050503, 0x05050507, 0x05050701, 0x0505070f, 0x05050903, 0x05050b07, 0x05050b0f, 0x05050f03, 0x05050f09, 0x05070101, 0x05070105, 0x0507010b, 0x05070303, 0x05070505, 0x05070509, 0x05070703, 0x05070707, 0x05070905, 0x05070b01, 0x05070d0d, 0x05090103, 0x0509010f, 0x05090501, 0x05090507, 0x05090705, 0x0509070b, 0x05090903, 0x05090f05, 0x05090f0b, 0x050b0109, 0x050b0303, 0x050b0505, 0x050b070f, 0x050b0901, 0x050b0b07, 0x050b0f01, 0x050d0101, 0x050d0105, 0x050d010f, 0x050d0503, 0x050d0b0b, 0x050d0d03, 0x050f010b, 0x050f0303, 0x050f050d, 0x050f0701, 0x050f0907, 0x050f0b01, 0x07010105, 0x07010303, 0x07010307, 0x0701030b, 0x0701030f, 0x07010505, 0x07010703, 0x07010707, 0x0701070b, 0x07010905, 0x07010909, 0x0701090f, 0x07010b03, 0x07010d07, 0x07010f03, 0x07030103, 0x07030107, 0x0703010b, 0x07030309, 0x07030503, 0x07030507, 0x07030901, 0x07030d01, 0x07030f05, 0x07030f0d, 0x07050101, 0x07050305, 0x07050501, 0x07050705, 0x07050709, 0x07050b01, 0x07070103, 0x07070301, 0x07070309, 0x07070503, 0x07070507, 0x0707050f, 0x07070701, 0x07070903, 0x07070907, 0x0707090f, 0x07070b0b, 0x07070f07, 0x07090107, 0x07090303, 0x0709030d, 0x07090505, 0x07090703, 0x07090b05, 0x07090d01, 0x07090d09, 0x070b0103, 0x070b0301, 0x070b0305, 0x070b050b, 0x070b0705, 0x070b0909, 0x070b0b0d, 0x070b0f07, 0x070d030d, 0x070d0903, 0x070f0103, 0x070f0107, 0x070f0501, 0x070f0505, 0x070f070b, 0x09010101, 0x09010109, 0x09010305, 0x09010501, 0x09010509, 0x0901050f, 0x09010705, 0x09010903, 0x09010b01, 0x09010f01, 0x09030105, 0x0903010f, 0x09030303, 0x09030307, 0x09030505, 0x09030701, 0x0903070b, 0x09030907, 0x09030b03, 0x09030b0b, 0x09050103, 0x09050107, 0x09050301, 0x0905030b, 0x09050503, 0x09050707, 0x09050901, 0x09050b0f, 0x09050d05, 0x09050f01, 0x09070109, 0x09070303, 0x09070307, 0x09070501, 0x09070505, 0x09070703, 0x0907070b, 0x09090101, 0x09090105, 0x09090509, 0x0909070f, 0x09090901, 0x09090f03, 0x090b010b, 0x090b010f, 0x090b0503, 0x090b0d05, 0x090d0307, 0x090d0709, 0x090d0d01, 0x090f0301, 0x090f030b, 0x090f0701, 0x090f0907, 0x090f0b03, 0x0b010105, 0x0b010301, 0x0b010309, 0x0b010505, 0x0b010901, 0x0b010909, 0x0b01090f, 0x0b010b05, 0x0b010d0d, 0x0b010f09, 0x0b030103, 0x0b030107, 0x0b03010b, 0x0b030305, 0x0b030503, 0x0b030705, 0x0b030f05, 0x0b050101, 0x0b050303, 0x0b050507, 0x0b050701, 0x0b05070d, 0x0b050b07, 0x0b070105, 0x0b07010f, 0x0b070301, 0x0b07050f, 0x0b070909, 0x0b070b03, 0x0b070d0b, 0x0b070f07, 0x0b090103, 0x0b090109, 0x0b090501, 0x0b090705, 0x0b09090d, 0x0b0b0305, 0x0b0b050d, 0x0b0b0b03, 0x0b0b0b07, 0x0b0d0905, 0x0b0f0105, 0x0b0f0109, 0x0b0f0505, 0x0d010303, 0x0d010307, 0x0d01030b, 0x0d010703, 0x0d010707, 0x0d010d01, 0x0d030101, 0x0d030501, 0x0d03050f, 0x0d030d09, 0x0d050305, 0x0d050709, 0x0d050905, 0x0d050b0b, 0x0d050d05, 0x0d050f01, 0x0d070101, 0x0d070309, 0x0d070503, 0x0d070901, 0x0d09050b, 0x0d090907, 0x0d090d05, 0x0d0b0101, 0x0d0b0107, 0x0d0b0709, 0x0d0b0d01, 0x0d0d010b, 0x0d0d0901, 0x0d0f0303, 0x0d0f0307, 0x0f010101, 0x0f010109, 0x0f01010f, 0x0f010501, 0x0f010505, 0x0f01070d, 0x0f010901, 0x0f010b09, 0x0f010d05, 0x0f030105, 0x0f030303, 0x0f030509, 0x0f030907, 0x0f03090b, 0x0f050103, 0x0f050109, 0x0f050301, 0x0f05030d, 0x0f050503, 0x0f050701, 0x0f050b03, 0x0f070105, 0x0f070705, 0x0f07070b, 0x0f070b07, 0x0f090103, 0x0f09010b, 0x0f090307, 0x0f090501, 0x0f090b01, 0x0f0b0505, 0x0f0b0905, 0x0f0d0105, 0x0f0d0703, 0x0f0f0101,}; diff --git a/edge0-strata/metal/src/engine/msl/mtp.metal b/edge0-strata/metal/src/engine/msl/mtp.metal new file mode 100644 index 00000000..7159330f --- /dev/null +++ b/edge0-strata/metal/src/engine/msl/mtp.metal @@ -0,0 +1,68 @@ +// metal/src/engine/msl/mtp.metal - the small steps of the MTP draft layer's input branches (include/strata/core/mtp.hpp) +// around mac-kern's kernels (GEMVs, the hyper-connection read, QSA, MoE, the head): +// +// mtp_rms y[r] = x[r] * rsqrt(mean(x[r]^2) + eps) * w rows of width n (w: (1 + w) folded by the host) +// mtp_xprep xprep (units.metal) of rows read with a stride: the prepared activation of one stream of rows +// mtp_combine R[t][c][d] = h[c][t][d] + e[t][d] the streams of fc_hidden plus fc_embedding +// +// Prepared activations: [4 kinds][n_tok][n] then S16 [n_tok][n / 16] (units.metal's xprep layout). + +#include +using namespace metal; + +struct RmsArgs { uint n; float eps; }; + +// grid (rows), 256 threads +kernel void mtp_rms(device const float* x [[buffer(0)]], device const float* w [[buffer(1)]], device float* y [[buffer(2)]], + constant RmsArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + device const float* xr = x + tg.x * a.n; + float ss = 0.0f; + for (uint d = tid; d < a.n; d += 256) ss = fma(xr[d], xr[d], ss); + ss = simd_sum(ss); + if (lane == 0) part[sg] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tot = 0.0f; + for (int i = 0; i < 8; ++i) tot += part[i]; + const float rs = rsqrt(tot / float(a.n) + a.eps); + device float* yr = y + tg.x * a.n; + for (uint d = tid; d < a.n; d += 256) yr[d] = xr[d] * rs * w[d]; +} + +struct XpArgs { uint n; uint n_tok; uint plane_tok; uint stride; }; + +// one thread per 16 values of one row: row t at x + t * stride; out planes of plane_tok rows +kernel void mtp_xprep(device const float* x [[buffer(0)]], device float* xp [[buffer(1)]], constant XpArgs& a [[buffer(2)]], + uint gid [[thread_position_in_grid]]) { + const uint per = a.n / 16; + if (gid >= per * a.n_tok) return; + const uint t = gid / per, i = gid % per; + device const float* src = x + t * a.stride + 16 * i; + const uint plane = a.plane_tok * a.n; + device float* raw = xp + t * a.n + 16 * i; + float v[16]; + for (int j = 0; j < 16; ++j) v[j] = src[j]; + float s0 = 0, s1 = 0; + for (int j = 0; j < 16; j += 2) { s0 += v[j]; s1 += v[j + 1]; } + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; + const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + for (int j = 0; j < 16; ++j) { + raw[j] = v[j]; + raw[plane + j] = v[j] * q2s[j & 7]; + raw[2 * plane + j] = v[j] * p4s[j & 3]; + raw[3 * plane + j] = (j & 1) ? v[j] * 0.00390625f : v[j]; + } + xp[4 * plane + t * per + i] = s0 + s1; +} + +struct CombArgs { uint n; uint hc; uint n_tok; uint h_stride; }; // h: stream c's rows at h + c * h_stride + +kernel void mtp_combine(device const float* h [[buffer(0)]], device const float* e [[buffer(1)]], device float* R [[buffer(2)]], + constant CombArgs& a [[buffer(3)]], uint gid [[thread_position_in_grid]]) { + const uint per = a.hc * a.n; + if (gid >= a.n_tok * per) return; + const uint t = gid / per, c = (gid % per) / a.n, d = gid % a.n; + R[gid] = h[c * a.h_stride + t * a.n + d] + e[t * a.n + d]; +} diff --git a/edge0-strata/metal/src/engine/msl/mtp_chain.metal b/edge0-strata/metal/src/engine/msl/mtp_chain.metal new file mode 100644 index 00000000..62500df2 --- /dev/null +++ b/edge0-strata/metal/src/engine/msl/mtp_chain.metal @@ -0,0 +1,92 @@ +// metal/src/engine/msl/mtp_chain.metal - the two steps that keep a whole MTP draft round on the GPU (one command +// buffer for every draft step; compiled after mac-kern's quant.metal and units.metal): +// +// mtp_pick draft j: the pick's token id (through the draft vocabulary) and its probability, the softmax of +// the draft head's logits at the pick (the host computed it after every step before) +// mtp_embed_F the next step's input: the picked token's embedding row dequantized from the model's token_embd +// (format F) - the host dequantized it before +// +// mtp_embed reads a row with mac-kern's unit decoders (units.metal Dec): a unit's dot with x is +// d1 * sum_{j < UNIT/2} q[j] xs[j] + d2 * sum_{j >= UNIT/2} q[j] xs[j] - (m1 xsum0 + m2 xsum1) +// with xs the prepared x (raw x times the kind's 2^-k) and xsum the raw sums of 16, so weight j of the unit is +// w = d * (q[j] * 2^-k) - m (d = d1 or d2 by half, m = m1 for the first 16 x values, m2 for the rest; +// a 16-wide unit has one sum: m1) +// and its position in the row is what X::load reads when x holds the positions 0 .. n-1. + +struct PickArgs { + uint nv; // the draft head's rows (the draft vocabulary, or the whole vocabulary) + uint j; // the draft step + uint use_map; // 1: map the head's row through the draft vocabulary +}; + +// one threadgroup of 1024 +kernel void mtp_pick(device const float* lg [[buffer(0)]], device const int* pick [[buffer(1)]], + device const int* vmap [[buffer(2)]], device int* dout [[buffer(3)]], device float* dprob [[buffer(4)]], + constant PickArgs& a [[buffer(5)]], ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[32]; + const int di = pick[0]; + const float mx = lg[di]; + float s = 0.0f; + for (uint i = tid; i < a.nv; i += 1024) s += exp(lg[i] - mx); + s = simd_sum(s); + if (lane == 0) part[sg] = s; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float t = 0.0f; + for (int k = 0; k < 32; ++k) t += part[k]; + dout[a.j] = a.use_map ? vmap[di] : di; + dprob[a.j] = 1.0f / t; + } +} + +struct EmbArgs { + uint n; // the row's width + uint row_bytes; + uint j; // dout[j] is the token +}; + +template static inline float kind_scale(int pos) { + const int k = XKind::K; + if (k == XK_Q2) { + const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f, 0.00006103515625f}; + return q2s[pos & 7]; + } + if (k == XK_P4) { + const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f}; + return p4s[pos & 3]; + } + if (k == XK_P2) return (pos & 1) ? 0.00390625f : 1.0f; + return 1.0f; +} + +// one threadgroup of 256; posv = the floats 0 .. n-1 +template +kernel void mtp_embed(device const uchar* W [[buffer(0)]], device const int* dout [[buffer(1)]], + device const float* posv [[buffer(2)]], device float* y [[buffer(3)]], constant EmbArgs& a [[buffer(4)]], + ushort tid [[thread_index_in_threadgroup]]) { + constexpr int UNIT = fmt::UNIT, H = UNIT / 2; + threadgroup uint tab[fmt::TAB]; + load_tab(tab, tid, 256); + threadgroup_barrier(mem_flags::mem_threadgroup); + device const uchar* row = W + (ulong) dout[a.j] * a.row_bytes; + const int U = int(a.n) / UNIT; + for (int u = tid; u < U; u += 256) { + float q[UNIT], ps[UNIT], xsum[2]; + float4 sc; + Dec::run(row, u, q, sc, tab); + X::load(posv, posv, u, ps, xsum); + for (int j = 0; j < UNIT; ++j) { + const int pos = int(ps[j]); + const float d = j < H ? sc.x : sc.y; + const float m = (UNIT == 16 || j < 16) ? sc.z : sc.w; + y[pos] = d * (q[j] * kind_scale(pos)) - m; + } + } +} + +// host names by the format's GGUF type number (mtp_embed_11 = Q3_K): EMB expands F before EMB_ stringifies it +#define EMB_(F) template [[host_name("mtp_embed_" #F)]] kernel decltype(mtp_embed) mtp_embed; +#define EMB(F) EMB_(F) +EMB(T_BF16) EMB(T_Q8_0) EMB(T_Q4_0) EMB(T_Q5_0) EMB(T_IQ4_NL) EMB(T_IQ4_XS) EMB(T_Q2_0) EMB(T_IQ3_S) EMB(T_Q3_K) +EMB(T_Q4_K) EMB(T_Q5_K) EMB(T_Q6_K) EMB(T_IQ3_XXS) diff --git a/edge0-strata/metal/src/engine/msl/ple.metal b/edge0-strata/metal/src/engine/msl/ple.metal new file mode 100644 index 00000000..e9e7b3c8 --- /dev/null +++ b/edge0-strata/metal/src/engine/msl/ple.metal @@ -0,0 +1,103 @@ +// metal/src/engine/msl/ple.metal - the PLE block (include/strata/kernels/ple.hpp) after its two projections, for a +// window of T rows, on the GPU (the CUDA engine's ple.cu math; its block sums are double there, float here): +// +// key = grouped_norm(pkey, norm_key) per stream c of 4: x * rsqrt(mean(x^2) + eps) * w +// query = grouped_norm(R, norm_query) +// gate[c] = sigmoid(sign(s) sqrt(max(|s|, 1e-6))), s = key[c] . query[c] / sqrt(n_embd) +// gated = value[d] * gate[c]; normalized = grouped_norm(gated, norm_conv) +// conv[c] = silu(sum_k kW[k + 4c] * row(t - (3 - k) * 3)[c]) rows: the history, then the window's own +// R += gated + conv +// +// History: hist[row + 9 * channel] (row-fastest, oldest first, the committed rows); the window's normalized rows +// are written to norm_w[t][channel] (the host appends the kept ones at commit). + +#include +using namespace metal; + +struct PleArgs { + uint n_embd; // 2560 + uint hc; // 4 + float eps; +}; + +static inline float tg_sum(float v, threadgroup float* part, ushort sg, ushort lane) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + float t = 0.0f; + for (int i = 0; i < 8; ++i) t += part[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return t; +} + +// grid (hc, T), 256 threads: y[t][c] = grouped_norm(x[t][c]) with w[c] +kernel void ple_gnorm(device const float* x [[buffer(0)]], device const float* w [[buffer(1)]], device float* y [[buffer(2)]], + constant PleArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd; + device const float* xc = x + (t * a.hc + c) * n; + float ss = 0.0f; + for (uint d = tid; d < n; d += 256) ss += xc[d] * xc[d]; + const float mean = tg_sum(ss, part, sg, lane) / float(n); + const float scale = 1.0f / sqrt(mean + a.eps); + device float* yc = y + (t * a.hc + c) * n; + for (uint d = tid; d < n; d += 256) yc[d] = xc[d] * scale * w[c * n + d]; +} + +// grid (hc, T), 256 threads: the gate, then gated = value * gate and normalized = grouped_norm(gated, norm_conv) +kernel void ple_gate_norm(device const float* key [[buffer(0)]], device const float* query [[buffer(1)]], + device const float* value [[buffer(2)]], device const float* w_conv [[buffer(3)]], + device float* gated [[buffer(4)]], device float* norm_w [[buffer(5)]], + constant PleArgs& a [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]]) { + threadgroup float part[8]; + const uint c = tg.x, t = tg.y, n = a.n_embd, row = (t * a.hc + c) * n; + float acc = 0.0f; + for (uint d = tid; d < n; d += 256) acc += key[row + d] * query[row + d]; + const float s = tg_sum(acc, part, sg, lane) * (1.0f / sqrt(float(n))); + const float mag = sqrt(max(fabs(s), 1e-6f)); + const float sgn = s > 0.0f ? 1.0f : (s < 0.0f ? -1.0f : 0.0f); + const float gate = 1.0f / (1.0f + exp(-(sgn * mag))); + float ss = 0.0f; + for (uint d = tid; d < n; d += 256) { + const float g = value[t * n + d] * gate; + gated[row + d] = g; + ss += g * g; + } + const float mean = tg_sum(ss, part, sg, lane) / float(n); + const float scale = 1.0f / sqrt(mean + a.eps); + for (uint d = tid; d < n; d += 256) norm_w[row + d] = gated[row + d] * scale * w_conv[c * n + d]; +} + +// grid (hc * n_embd / 256, T), 256 threads: the conv over [history | the window's rows], SiLU, and the residual +// PLE_ROWS (function constant 32): the rows are different sequences (batch slots): row t convolves its slot's history only (slot +// stride rinfo[7] floats, row t at rinfo[8 + 4 t] = {slot, ...}); the three taps before the row itself always read the history. +constant bool PLE_ROWS_C [[function_constant(32)]]; +constant bool ple_rows = is_function_constant_defined(PLE_ROWS_C) && PLE_ROWS_C; + +kernel void ple_conv_add(device const float* hist [[buffer(0)]], device const float* norm_w [[buffer(1)]], + device const half* kw [[buffer(2)]], device const float* gated [[buffer(3)]], + device float* R [[buffer(4)]], constant PleArgs& a [[buffer(5)]], + device const uint* rinfo [[buffer(6), function_constant(ple_rows)]], + uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) { + const uint ch = tg.x * 256 + tid, t = tg.y, hcd = a.hc * a.n_embd; + if (ch >= hcd) return; + if (ple_rows) hist += ulong(rinfo[8 + 4 * t]) * rinfo[7]; + float acc = 0.0f; + for (int k = 0; k < 4; ++k) { + float v; + if (ple_rows) { + v = k == 3 ? norm_w[t * hcd + ch] : hist[uint(9 - (3 - k) * 3) + 9 * ch]; + } else { + const int p = int(t) - (3 - k) * 3; // the row this tap reads: < 0 = the committed history + v = p >= 0 ? norm_w[uint(p) * hcd + ch] : hist[uint(9 + p) + 9 * ch]; + } + acc += float(kw[k + 4 * ch]) * v; + } + const float conv = acc / (1.0f + exp(-acc)); + device float* r = R + t * hcd; + r[ch] = r[ch] + gated[t * hcd + ch] + conv; +} diff --git a/edge0-strata/metal/src/engine/prompt_chain.inc b/edge0-strata/metal/src/engine/prompt_chain.inc new file mode 100644 index 00000000..1278edb8 --- /dev/null +++ b/edge0-strata/metal/src/engine/prompt_chain.inc @@ -0,0 +1,571 @@ +// metal/src/engine/prompt_chain.inc (included by window.cpp) - the prompt through mac-kern's prompt kernels +// (metal/tests/prompt_chain.cpp's chain, 6f6c555b), into the decode window's own state: chunks of C rows, every +// dense matrix dequantized to fp16 and multiplied on the matrix units, the routed experts as grouped GEMMs, the GDN +// state by the parallel prompt scan, QSA K/V and indexer keys into the window's caches. Added for the engine: +// the PLE block after layer 0 (GEMMs + msl/ple.metal), every row's final residual (the drafter's prompt rows), the +// window's history parity, RoPE table and pending fold. No head: the last prompt token runs as a decode window. +// +// The 8-row-window prompt path stays as the fallback: STRATA_METAL_PROMPT=window, and prompts shorter than +// STRATA_METAL_PROMPT_MIN rows (default 32). STRATA_METAL_PROMPT_CHUNK = C (default 2048). The kernel variants +// follow prompt_chain's defaults: matrix-unit register operands where the GPU has the M5's right-input layout +// (COOP_K4, function constant 7), else the older expert and attention kernels (the M4). + +namespace { +constexpr int PC_QB = 2051, PC_QCHUNKS = 33, PC_STAGE_GU_MIN = 2048; +struct PcDqArgs { uint32_t n_in, n_out, row_bytes, perm; }; +struct PcGemmArgs { int M, N, K, pad; }; +struct PcEpiArgs { uint32_t n_tok, n_embd, hc, hc_lr; }; +struct PcGdnPArgs { uint32_t n_tok, h_k, h_v, channels; float eps; uint32_t f32; }; +struct PcMmPlanArgs { uint32_t n_ent, topk, n_expert, tm, n_blocks; }; +struct PcMmArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff; }; +struct PcLatentArgs { uint32_t tokens, groups, hidden, width, topk, tile_rows; }; +int pc_unit_of(int t) { + return t == ref::Q2_0 || t == ref::IQ3_S || t == ref::BF16 || t == ref::F32 || t == 18 || t == 16 || t == 17 || t == 22 ? 16 : 32; +} +} // namespace + +bool Window::coop_k4() { + if (const char* e = std::getenv("COOP_K4")) return std::atoi(e) != 0; + const int lib = cx_->compile({"mm_bench.metal"}, "#define NO_GGML_COMMON 1\n"); + mtl::Pipeline p = cx_->pipeline(lib, "coop_layout_16_64_32"); + mtl::Buffer out = cx_->alloc(4 * 3 * 32 * 128 * 4); + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(out)}; + d.tg[0] = 32; + cx_->run({d}); + const int* o = out.as(); + bool k4 = true; + for (int l = 0; l < 32 && k4; ++l) { + const int* q = o + l * 128 * 4; + const int cap = q[2]; + for (int i = 0; i + 3 < cap && k4; i += 4) + for (int j = 1; j < 4; ++j) + if (q[(i + j) * 4] != q[i * 4] + j || q[(i + j) * 4 + 1] != q[i * 4 + 1]) { k4 = false; break; } + } + cx_->free(out); + return k4; +} + +bool Window::pc_setup(int C) { + PromptBufs& b = pb_; + if (b.C >= C) return true; + const double t0 = now_ms(); + const char* tune = std::getenv("STRATA_METAL_LATENT_PROMPT_TUNE"); + const bool tuned = latent_on_ && tune && std::strcmp(tune, "1") == 0; + if (b.C == 0) { + std::vector kernels{"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", + "gr.metal", "gemv.metal", "prompt.metal"}; + if (latent_on_) kernels.push_back("latent_prompt.metal"); + if (tuned) kernels.push_back("latent_prompt_q3_stage.metal"); + b.lib = cx_->compile(kernels); + b.lgdn = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + b.lqsa = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + b.k4 = coop_k4(); + } + b = PromptBufs{b.lib, b.lgdn, b.lqsa, b.k4}; // (re)allocated below for the larger chunk + b.v3e = !latent_on_ && b.k4 && T("blk.0.ffn_gate_exps.weight").type < 144; + b.C = C; + const size_t MT = (size_t) C; + const int MB = (int) max_cells_ / 4 + 2; + b.expert_tm = tuned ? 32 : 16; + b.maxrows = (int) MT * TOPK + NEXP * b.expert_tm; + const int nblk_max = ((int) MT * TOPK + 1023) / 1024; + auto F32 = [&](size_t n) { return cx_->alloc(4 * n); }; + auto F16 = [&](size_t n) { return cx_->alloc(2 * n); }; + b.R = F32(MT * HC * N); b.bo = F32(MT * N); b.inj[0] = F32(MT * HC); b.inj[1] = F32(MT * HC); b.xn = F32(MT * HC * N); + b.xnh = F16(MT * HC * N); b.V = F32(MT * (LR + HC)); b.loh = F16(MT * LR); b.mixed = F32(MT * N); b.mixh = F16(MT * N); + b.fused = gdn_fused_inputs(cx_->device_name()); + if (b.fused) { // the four GDN input projections as one GEMM: [qkv | z | alpha | beta] rows of 16480 floats + b.yf = F32(MT * 16480); + b.qkv = F32(1); b.z = F32(1); b.al = F32(1); b.be = F32(1); + } else { + b.qkv = F32(MT * 10240); b.z = F32(MT * 6144); b.al = F32(MT * 48); b.be = F32(MT * 48); + } + b.yg = F32(MT * 6144); + b.ygh = F16(MT * 6144); b.gqn = F32(MT * 16 * 128); b.gkn = F32(MT * 16 * 128); b.gvc = F32(MT * 6144); + b.ggb = F32(MT * 48 * 2); b.gkq = F32(MT * 16); b.go = F32(MT * 6144); + b.qfull = F32(MT * 12288); b.kcur = F32(MT * 512); b.vcur = F32(MT * 512); b.iqr = F32(MT * 512); b.ikr = F32(MT * 128); + b.qcur = F32(MT * 24 * 256); b.qcurh = F16(MT * 24 * 256); b.qidx = F32(MT * 4 * 128); b.scores = F32(MT * MB); + b.qids = cx_->alloc(4 * MT * PC_QB); b.part = F32(MT * PC_QCHUNKS * 24 * 258); b.XQ = F32(4 * MT * 6144 + MT * 384); + b.xqh = F16(MT * 6144); + b.logits = F32(MT * NEXP); b.ids = cx_->alloc(4 * MT * TOPK); b.wts = F32(MT * TOPK); + b.bcnt = cx_->alloc(4 * (size_t) nblk_max * NEXP); b.boff = cx_->alloc(4 * (size_t) nblk_max * NEXP); + b.eoff = cx_->alloc(4 * NEXP); b.tiles = cx_->alloc(8 * (size_t) (b.maxrows / b.expert_tm + 1)); b.meta = cx_->alloc(64); + b.stok = cx_->alloc(4 * (size_t) b.maxrows); b.sw = F32(b.maxrows); b.pos = cx_->alloc(4 * MT * TOPK); + const int expert_width = latent_on_ ? (int) latent_args_.latent : N; + b.XS = F16((size_t) b.maxrows * expert_width); b.HS = F16((size_t) b.maxrows * NFF); + b.YE = F32((size_t) b.maxrows * expert_width); + if (latent_on_) { + const size_t groups = latent_args_.groups; + b.lat_w = F16(groups * expert_width * N); + b.lat_proj = F32(groups * MT * N); + b.lat_pooled = F16(groups * MT * expert_width); + // One layer of gate/up operands is reused across chunks and layers. + if (tuned) b.lat_gu16 = F16(2ull * NEXP * NFF * expert_width); + } + b.yr = F32(MT * N); b.gsh = F32(MT * NFF); b.ush = F32(MT * NFF); b.hh = F16(MT * NFF); b.ys = F32(MT * N); b.sgv = F32(MT); + b.WD = F16((size_t) (LR + HC) * HC * N); b.WU = F16((size_t) HC * N * LR); + if (b.fused) { + b.Wcat = F16((size_t) 16480 * N); + b.Wqkv = F16(1); b.Wz = F16(1); b.Wal = F16(1); b.Wbe = F16(1); + } else { + b.Wqkv = F16((size_t) 10240 * N); b.Wz = F16((size_t) 6144 * N); b.Wal = F16((size_t) 48 * N); b.Wbe = F16((size_t) 48 * N); + } + b.Wout = F16((size_t) N * 6144); + b.Wq = F16((size_t) 12288 * N); b.Wk = F16((size_t) 512 * N); b.Wv = F16((size_t) 512 * N); b.Wiq = F16((size_t) 512 * N); + b.Wik = F16((size_t) 128 * N); b.Wr = F16((size_t) NEXP * N); b.Wsg = F16((size_t) N); b.Wshg = F16((size_t) NFF * N); + b.Wshu = F16((size_t) NFF * N); b.Wshd = F16((size_t) N * NFF); + // the PLE block's rows and projections + b.pemb = F32(MT * N); b.pembh = F16(MT * N); b.pkey = F32(MT * 10240); b.pval = F32(MT * N); b.keyn = F32(MT * 10240); + b.qn = F32(MT * 10240); b.gated = F32(MT * 10240); b.normW = F32(MT * 10240); + if (had_on_) b.mixr = F32(MT * N); + b.Wpk = F16((size_t) 10240 * N); b.Wpv = F16((size_t) N * N); + // the expert GEMM tiles (prompt_chain's defaults): v3 gate/up 16,32,32,4 and down 16,64,32,2 with the M5 layout, + // else v2 + b.gu3[0] = 16; b.gu3[1] = 32; b.gu3[2] = 32; b.gu3[3] = 4; + b.dn3[0] = 16; b.dn3[1] = 64; b.dn3[2] = 32; b.dn3[3] = 2; + b.meta.as()[2] = b.v3e ? NFF / (b.gu3[1] * b.gu3[3]) : NFF / 64; + b.meta.as()[5] = b.v3e ? N / (b.dn3[1] * b.dn3[3]) : expert_width / 64; + std::fprintf(stderr, "strata metal: prompt kernels (mac-kern's prompt chain) for chunks of %d rows, %s, ready in %.0f ms\n", + C, b.k4 ? "matrix-unit register operands (COOP_K4)" : "the older expert / attention kernels (no M5 layout)", + now_ms() - t0); + return true; +} + +// one chunk of nT rows at cells pos0.. (the window state is exactly at pos0 before, at pos0 + nT after) +bool Window::pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& err) { + PromptBufs& b = pb_; + strata::core::SessionState& ss = *ss_; + const double th0 = now_ms(); + // ---- the rows: the embedding into every stream; the PLE table rows + const gguf::Tensor& te = T("token_embd.weight"); + const size_t rb = ref::row_bytes(te.type, N); + mrope_sync(pos0, nT); + for (int t = 0; t < nT; ++t) { + float* r = b.R.as() + (size_t) t * HC * N; + if (const float* img = img_row(pos0 + t)) std::memcpy(r, img, (size_t) N * 4); + else ref::dequant(te.type, (const uint8_t*) W_.ptr + te.offset + (size_t) tokens[t] * rb, r, N); + for (int c = 1; c < HC; ++c) std::memcpy(r + (size_t) c * N, r, (size_t) N * 4); + } + if (ple_on_) { + std::vector rows((size_t) nT * strata::kernels::PLE_N_HEADS); + int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]}; + for (int t = 0; t < nT; ++t) { + const int32_t tk = (int32_t) tokens[t]; + strata::kernels::ngram_rows(&tk, prev, 1, ss.ple.consts, rows.data() + (size_t) t * strata::kernels::PLE_N_HEADS); + prev[0] = prev[1]; + prev[1] = tk; + } + if (!ss.ple.table->gather_batch(rows.data(), (size_t) nT, b.pemb.as(), err)) return false; + } + b.host_ms += now_ms() - th0; + std::vector ds; + ds.reserve(2400); + static const bool stage_profile = std::getenv("STRATA_METAL_LATENT_PROMPT_PROFILE") != nullptr; + static const bool latent_qsa_mm = [] { + const char* e = std::getenv("STRATA_METAL_LATENT_PROMPT_QSA_MM"); + return e && std::strcmp(e, "1") == 0; + }(); + const bool qsa_mm = b.k4 || (latent_on_ && latent_qsa_mm); + std::vector> stage_ends; + auto end_stage = [&](const char* name) { + if (stage_profile && latent_on_) stage_ends.emplace_back(name, ds.size()); + }; + nb_ = false; + if (pend_fold_ > 0) { // the last window's commit, first + fold(ds, pend_fold_, 1 - hp_); + pend_fold_ = 0; + } + auto P_ = [&](int l, const std::string& fn, const mtl::Constants& c = {}, const std::string& key = "") { + return P(l, fn, c, "pc:" + (key.empty() ? fn : key)); + }; + auto D = [&](const mtl::Pipeline* p, std::vector a, uint32_t gx, uint32_t gy, uint32_t tg, uint32_t gz = 1) { + mtl::Dispatch d; + d.p = p; + d.args = std::move(a); + d.grid[0] = gx; d.grid[1] = gy; d.grid[2] = gz; + d.tg[0] = tg; + mtl::require_fit(d, "prompt-chain dispatch"); + ds.push_back(std::move(d)); + }; + auto gemm = [&](const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Y, int M, int Nn, int K, + size_t a_off = 0, size_t b_off = 0, size_t y_off = 0) { + int tm = 64, tn = Nn <= 64 ? 32 : 64; + gemm_tile(cx_->device_name(), Nn, K, tm, tn); + const std::string fn = "gemm_f16_" + std::to_string(tm) + "_" + std::to_string(tn) + "_4"; + PcGemmArgs g{M, Nn, K, 0}; + D(P_(b.lib, fn), {mtl::Arg::B(A, a_off), mtl::Arg::B(B, b_off), mtl::Arg::B(Y, y_off), mtl::Arg::V(g)}, + (Nn + tn - 1) / tn, (M + tm - 1) / tm, 128); + }; + auto dq = [&](const gguf::Tensor& t, const mtl::Buffer& out, size_t out_off = 0, uint32_t perm = 0) { + const int n_in = (int) t.shape[0], n_out = t.shape.size() > 1 ? (int) t.shape[1] : 1; + PcDqArgs a{(uint32_t) n_in, (uint32_t) n_out, (uint32_t) ref::row_bytes(t.type, n_in), perm}; + D(P_(b.lib, "dq_f16_" + std::to_string(t.type)), {mtl::Arg::B(W_, t.offset), mtl::Arg::B(out, out_off), mtl::Arg::V(a)}, + (n_in / pc_unit_of(t.type) + 63) / 64, n_out, 64); + }; + auto act16 = [&](const mtl::Buffer& x, const mtl::Buffer& xh, size_t n) { + const uint32_t nn = (uint32_t) n; + D(P_(b.lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 255) / 256, 1, 256); + }; + const PcLatentArgs la{(uint32_t) nT, latent_args_.groups, N, latent_args_.latent, TOPK, b.expert_tm}; + auto latent_project = [&](size_t offset, uint32_t in, uint32_t out, uint32_t row_bytes, + const mtl::Buffer& input, bool group_inputs) { + const PcDqArgs args{in, la.groups * out, row_bytes, 0}; + D(P_(b.lib, "dq_f16_2"), {mtl::Arg::B(latent_du_, offset), mtl::Arg::B(b.lat_w), mtl::Arg::V(args)}, + (in / 32 + 63) / 64, args.n_out, 64); + for (uint32_t group = 0; group < la.groups; ++group) + gemm(input, b.lat_w, b.lat_proj, nT, out, in, group_inputs ? 2ull * group * nT * in : 0, + 2ull * group * out * in, 4ull * group * nT * out); + }; + int cur_inj = 0; + auto hc_read = [&](const std::string& pre, bool apply) { + const gguf::Tensor &tn = T(pre + "_norm.weight"), &tdn = T(pre + "_down.weight"), &tin = T(pre + "_inject.weight"), + &tup = T(pre + "_up.weight"); + dq(tdn, b.WD); + dq(tin, b.WD, (size_t) 2 * LR * HC * N); + GrArgs a{N, HC, LR, (uint32_t) nT, apply ? 1u : 0u, eps_}; + PcEpiArgs ea{(uint32_t) nT, N, HC, LR}; + dq(tup, b.WU, 0, HC); + D(P_(b.lib, "gr_norm_h"), {mtl::Arg::B(b.R), mtl::Arg::B(b.bo), mtl::Arg::B(b.inj[cur_inj]), mtl::Arg::B(W_, tn.offset), + mtl::Arg::B(b.xnh), mtl::Arg::V(a)}, HC, nT, 256); + gemm(b.xnh, b.WD, b.V, nT, LR + HC, HC * N); + D(P_(b.lib, "gr_epi_down"), {mtl::Arg::B(b.V), mtl::Arg::B(b.loh), mtl::Arg::B(b.inj[1 - cur_inj]), mtl::Arg::V(ea)}, + (nT * (LR + HC) + 255) / 256, 1, 256); + PcGemmArgs g{nT, HC * N, LR, 0}; + D(P_(b.lib, "gemm_hc_up_64_64_4"), {mtl::Arg::B(b.loh), mtl::Arg::B(b.WU), mtl::Arg::B(b.xnh), mtl::Arg::B(b.mixed), + mtl::Arg::B(b.mixh), mtl::Arg::V(g), mtl::Arg::V(ea)}, + HC * N / 64, (nT + 63) / 64, 128); + cur_inj = 1 - cur_inj; + }; + const int hp = hp_; + int gi = 0, qi = 0; + for (int l = 0; l < 48; ++l) { + const std::string bl = "blk." + std::to_string(l) + "."; + hc_read(bl + "hc_attn", l > 0 && !(l == 1 && ple_on_)); + if (l % 4 != 3) { + dq(T(bl + "ssm_out.weight"), b.Wout); + PcGdnPArgs ga{(uint32_t) nT, 16, 48, 10240, eps_, 0}; + if (b.fused) { + dq(T(bl + "attn_qkv.weight"), b.Wcat, 0); + dq(T(bl + "attn_gate.weight"), b.Wcat, (size_t) 2 * 10240 * N); + dq(T(bl + "ssm_alpha.weight"), b.Wcat, (size_t) 2 * (10240 + 6144) * N); + dq(T(bl + "ssm_beta.weight"), b.Wcat, (size_t) 2 * (10240 + 6144 + 48) * N); + gemm(b.mixh, b.Wcat, b.yf, nT, 16480, N); + D(P_(b.lgdn, "gdn_pre", mtl::Constants().i(7, 16480), "gdn_pre_f"), + {mtl::Arg::B(b.yf), mtl::Arg::B(ghist_[hp][(size_t) gi]), mtl::Arg::B(W_, T(bl + "ssm_conv1d.weight").offset), + mtl::Arg::B(b.yf, (size_t) 4 * (10240 + 6144)), mtl::Arg::B(b.yf, (size_t) 4 * (10240 + 6144 + 48)), + mtl::Arg::B(W_, T(bl + "ssm_dt.bias").offset), mtl::Arg::B(W_, T(bl + "ssm_a").offset), mtl::Arg::B(b.gqn), + mtl::Arg::B(b.gkn), mtl::Arg::B(b.gvc), mtl::Arg::B(b.ggb), mtl::Arg::B(b.gkq), + mtl::Arg::B(ghist_[1 - hp][(size_t) gi]), mtl::Arg::V(ga)}, + 16 + 24, nT, 256); + } else { + dq(T(bl + "attn_qkv.weight"), b.Wqkv); + dq(T(bl + "attn_gate.weight"), b.Wz); + dq(T(bl + "ssm_alpha.weight"), b.Wal); + dq(T(bl + "ssm_beta.weight"), b.Wbe); + gemm(b.mixh, b.Wqkv, b.qkv, nT, 10240, N); + gemm(b.mixh, b.Wz, b.z, nT, 6144, N); + gemm(b.mixh, b.Wal, b.al, nT, 48, N); + gemm(b.mixh, b.Wbe, b.be, nT, 48, N); + D(P_(b.lgdn, "gdn_pre"), + {mtl::Arg::B(b.qkv), mtl::Arg::B(ghist_[hp][(size_t) gi]), mtl::Arg::B(W_, T(bl + "ssm_conv1d.weight").offset), + mtl::Arg::B(b.al), mtl::Arg::B(b.be), mtl::Arg::B(W_, T(bl + "ssm_dt.bias").offset), + mtl::Arg::B(W_, T(bl + "ssm_a").offset), mtl::Arg::B(b.gqn), mtl::Arg::B(b.gkn), mtl::Arg::B(b.gvc), + mtl::Arg::B(b.ggb), mtl::Arg::B(b.gkq), mtl::Arg::B(ghist_[1 - hp][(size_t) gi]), mtl::Arg::V(ga)}, + 16 + 24, nT, 256); + } + D(P_(b.lgdn, gdn_scan_name(cx_->device_name(), nT), mtl::Constants().i(6, 0)), + {mtl::Arg::B(gst_[(size_t) gi]), mtl::Arg::B(b.gqn), mtl::Arg::B(b.gkn), mtl::Arg::B(b.gvc), mtl::Arg::B(b.ggb), + mtl::Arg::B(b.gkq), mtl::Arg::B(b.go), mtl::Arg::V(ga)}, + 48 * 4, 1, 32 * 4); + if (b.fused) + D(P_(b.lgdn, "gdn_post", mtl::Constants().i(7, 16480), "gdn_post_f"), + {mtl::Arg::B(b.go), mtl::Arg::B(b.yf, (size_t) 4 * 10240), mtl::Arg::B(W_, T(bl + "ssm_norm.weight").offset), + mtl::Arg::B(b.ygh), mtl::Arg::B(b.yg), mtl::Arg::V(ga)}, + 48, nT, 128); + else + D(P_(b.lgdn, "gdn_post"), {mtl::Arg::B(b.go), mtl::Arg::B(b.z), mtl::Arg::B(W_, T(bl + "ssm_norm.weight").offset), + mtl::Arg::B(b.ygh), mtl::Arg::B(b.yg), mtl::Arg::V(ga)}, + 48, nT, 128); + gemm(b.ygh, b.Wout, b.bo, nT, N, 6144); + ++gi; + } else { + dq(T(bl + "attn_q.weight"), b.Wq); + dq(T(bl + "attn_k.weight"), b.Wk); + dq(T(bl + "attn_v.weight"), b.Wv); + dq(T(bl + "indexer.q_proj.weight"), b.Wiq); + dq(T(bl + "indexer.k_proj.weight"), b.Wik); + dq(T(bl + "attn_output.weight"), b.Wout); + gemm(b.mixh, b.Wq, b.qfull, nT, 12288, N); + gemm(b.mixh, b.Wk, b.kcur, nT, 512, N); + gemm(b.mixh, b.Wv, b.vcur, nT, 512, N); + gemm(b.mixh, b.Wiq, b.iqr, nT, 512, N); + gemm(b.mixh, b.Wik, b.ikr, nT, 128, N); + QsaArgs a{(uint32_t) pos0, 0, (uint32_t) max_cells_, (uint32_t) nT, eps_}; + D(P_(b.lqsa, "qsa_prep", qsa_c(), qsa_k()), + {mtl::Arg::B(b.qfull), mtl::Arg::B(b.kcur), mtl::Arg::B(b.vcur), mtl::Arg::B(b.iqr), mtl::Arg::B(b.ikr), + mtl::Arg::B(W_, T(bl + "attn_q_norm.weight").offset), mtl::Arg::B(W_, T(bl + "attn_k_norm.weight").offset), + mtl::Arg::B(W_, T(bl + "indexer.q_norm.weight").offset), mtl::Arg::B(W_, T(bl + "indexer.k_norm.weight").offset), + mtl::Arg::B(cs_), mtl::Arg::B(b.qcur), mtl::Arg::B(b.qidx), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), + mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(tail_[(size_t) qi]), + mtl::Arg::B(dead_[(size_t) qi]), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::V(a), mtl::Arg::B(mrope_)}, + 33, nT, 256); + if (pos0 + nT > PC_QB) { // some row's context is past the budget: the indexer selects + if (qsa_mm) + D(P_(b.lqsa, "qsa_score_mm"), {mtl::Arg::B(b.qidx), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), + mtl::Arg::B(b.scores), mtl::Arg::V(a)}, + (uint32_t) (((pos0 + nT) / 4 + 1 + 63) / 64), (nT + 15) / 16, 128); + else + D(P_(b.lqsa, "qsa_score"), {mtl::Arg::B(b.qidx), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), + mtl::Arg::B(b.scores), mtl::Arg::V(a)}, + (uint32_t) (((pos0 + nT) / 4 + 1 + 7) / 8), nT, 256); + } + D(P_(b.lqsa, "qsa_topk"), {mtl::Arg::B(b.scores), mtl::Arg::B(b.qids), mtl::Arg::V(a)}, nT, 1, 1024); + if (qsa_mm) { + const uint32_t nq = (uint32_t) nT * 24 * 256; + D(P_(b.lqsa, "qsa_q_f16"), {mtl::Arg::B(b.qcur), mtl::Arg::B(b.qcurh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 256); + D(P_(b.lqsa, "qsa_attn_mm", mtl::Constants().b(7, b.k4), b.k4 ? "qsa_attn_mm#k4" : "qsa_attn_mm#scalar"), + {mtl::Arg::B(b.qcurh), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), mtl::Arg::B(ksc_[(size_t) qi]), + mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(b.qids), mtl::Arg::B(b.part), mtl::Arg::V(a)}, + PC_QCHUNKS, 2, 32, nT); + } else { + D(P_(b.lqsa, "qsa_attn"), {mtl::Arg::B(b.qcur), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), + mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(b.qids), + mtl::Arg::B(b.part), mtl::Arg::V(a)}, + PC_QCHUNKS, 2, 256, nT); + } + D(P_(b.lqsa, "qsa_merge"), {mtl::Arg::B(b.part), mtl::Arg::B(b.qfull), mtl::Arg::B(b.XQ), mtl::Arg::V(a)}, 24, nT, 256); + act16(b.XQ, b.xqh, (size_t) nT * 6144); + gemm(b.xqh, b.Wout, b.bo, nT, N, 6144); + ++qi; + } + end_stage("attention"); + // ---- MoE + hc_read(bl + "hc_ffn", true); + dq(T(bl + "ffn_gate_inp.weight"), b.Wr); + dq(T(bl + "ffn_gate_inp_shexp.weight"), b.Wsg); + dq(T(bl + "ffn_gate_shexp.weight"), b.Wshg); + dq(T(bl + "ffn_up_shexp.weight"), b.Wshu); + dq(T(bl + "ffn_down_shexp.weight"), b.Wshd); + gemm(b.mixh, b.Wr, b.logits, nT, NEXP, N); + const uint32_t nt = (uint32_t) nT; + D(P_(b.lib, "route_topk"), {mtl::Arg::B(b.logits), mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::V(nt)}, (nT + 7) / 8, 1, 256); + const int nent = nT * TOPK, nblk = (nent + 1023) / 1024; + PcMmPlanArgs pa{(uint32_t) nent, TOPK, NEXP, b.expert_tm, (uint32_t) nblk}; + const uint32_t nclr = (uint32_t) b.maxrows; + D(P_(b.lib, "mm_clear"), {mtl::Arg::B(b.stok), mtl::Arg::V(nclr)}, (b.maxrows + 255) / 256, 1, 256); + D(P_(b.lib, "mm_hist"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024); + D(P_(b.lib, "mm_scan"), {mtl::Arg::B(b.bcnt), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.tiles), mtl::Arg::B(b.meta), + mtl::Arg::V(pa)}, 1, 1, NEXP); + D(P_(b.lib, "mm_place"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.stok), + mtl::Arg::B(b.sw), mtl::Arg::B(b.pos), mtl::Arg::V(pa)}, nblk, 1, 1024); + PcMmArgs ma{}; + int gu_type, down_type; + const mtl::Buffer* expert_weights = &W_; + size_t ebase; + const size_t labels = (size_t) l * NEXP * sizeof(int32_t); + size_t du_offset = 0; + if (latent_on_) { + const auto& a = latent_args_; + ma = {0, a.up_off, a.down_off, a.expert_bytes, a.expert_bytes, a.gu_row, a.down_row, a.latent, NFF}; + gu_type = ref::Q3_K; + down_type = ref::IQ4_NL; + expert_weights = &latent_a_; + ebase = (size_t) l * NEXP * a.expert_bytes; + du_offset = (size_t) l * (latent_dense_.d_bytes() + latent_dense_.u_bytes()); + latent_project(du_offset, N, la.width, (uint32_t) latent_dense_.d_row_bytes(), b.mixh, false); + } else { + const gguf::Tensor &eg = T(bl + "ffn_gate_exps.weight"), &eu = T(bl + "ffn_up_exps.weight"), + &ed = T(bl + "ffn_down_exps.weight"); + const auto gr = (uint32_t) ref::row_bytes(eg.type, N), dr = (uint32_t) ref::row_bytes(ed.type, NFF); + ma = {eg.offset, eu.offset, ed.offset, (uint64_t) gr * NFF, (uint64_t) dr * N, gr, dr, N, NFF}; + gu_type = eg.type; + down_type = ed.type; + ebase = moe_rebase(ma.gate_off, ma.up_off, ma.down_off); + } + const uint32_t ne = N; + const bool hg = !latent_on_ && had_gu(l), hd = !latent_on_ && had_dn(l); + if (hg) had_rows(ds, b.lib, "pc:", b.mixed, 0, b.mixr, 0, nT, N, hsx_.handle ? &hsx_ : nullptr, false); + mtl::Dispatch g; + if (latent_on_) { + g.p = P_(b.lib, "latent_prompt_gather"); + g.args = {mtl::Arg::B(b.lat_proj), mtl::Arg::B(latent_inv_, (size_t) l * la.width * sizeof(float)), + mtl::Arg::B(latent_labels_, labels), mtl::Arg::B(b.tiles), mtl::Arg::B(b.stok), + mtl::Arg::B(b.XS), mtl::Arg::V(la)}; + } else { + g.p = P_(b.lib, "mm_gather"); + g.args = {mtl::Arg::B(hg ? b.mixr : b.mixed), mtl::Arg::B(b.stok), mtl::Arg::B(b.XS), mtl::Arg::V(ne)}; + } + g.tg[0] = 256; + g.indirect = &b.meta; + g.ind_off = 8 * 4; + ds.push_back(g); + end_stage("route_D"); + auto tl3 = [](const int* t) { return std::to_string(t[0]) + "_" + std::to_string(t[1]) + "_" + std::to_string(t[2]) + "_" + std::to_string(t[3]); }; + mtl::Dispatch gu; + const bool staged = b.lat_gu16.handle && nT >= PC_STAGE_GU_MIN; + if (staged) { + const uint32_t q3[4] = {ma.n_embd, ma.n_ff, ma.gu_row, uint32_t(ma.gu_stride)}; + D(P_(b.lib, "latent_prompt_q3_stage"), + {mtl::Arg::B(*expert_weights, ebase), mtl::Arg::B(b.lat_gu16), mtl::Arg::V(q3)}, + (2 * ma.n_ff * (ma.n_embd / 32) + 127) / 128, NEXP, 128); + gu.p = P_(b.lib, "latent_prompt_gu_f16"); + } else if (b.v3e) { + const std::string fn = "mm_gu3_" + std::to_string(gu_type) + "_" + tl3(b.gu3); + gu.p = P_(b.lib, fn, mtl::Constants().b(7, true), fn + "#k4"); + } else { + const char* latent_tile = b.expert_tm == 32 ? "_32_64_32_2" : "_16_64_64_2"; + const std::string fn = "mm_gu2_" + std::to_string(gu_type) + (latent_on_ ? latent_tile : gu2_tile(gu_type)); + gu.p = P_(b.lib, fn, mtl::Constants().i(5, 0), fn + "#p0"); + } + gu.args = {staged ? mtl::Arg::B(b.lat_gu16) : mtl::Arg::B(*expert_weights, ebase), mtl::Arg::B(b.XS), mtl::Arg::B(b.HS), mtl::Arg::B(b.tiles), mtl::Arg::V(ma)}; + gu.tg[0] = b.v3e ? 32 * b.gu3[3] : 64; + gu.indirect = &b.meta; + gu.ind_off = 2 * 4; + mtl::require_fit(gu, "prompt gate/up dispatch"); + ds.push_back(gu); + if (hd) had_rows_ind(ds, b.lib, "pc:", b.HS, b.meta); + mtl::Dispatch dn; + if (b.v3e) { + const std::string fn = "mm_down3_" + std::to_string(down_type) + "_" + tl3(b.dn3); + dn.p = P_(b.lib, fn, mtl::Constants().b(7, true), fn + "#k4"); + } else { + const std::string tile = b.expert_tm == 32 ? "_32_64_32_2" : dn2_tile(down_type); + dn.p = P_(b.lib, "mm_down2_" + std::to_string(down_type) + tile); + } + dn.args = {mtl::Arg::B(*expert_weights, ebase), mtl::Arg::B(b.HS), mtl::Arg::B(b.YE), mtl::Arg::B(b.tiles), mtl::Arg::V(ma)}; + dn.tg[0] = b.v3e ? 32 * b.dn3[3] : 64; + dn.indirect = &b.meta; + dn.ind_off = 5 * 4; + mtl::require_fit(dn, "prompt down dispatch"); + ds.push_back(dn); + end_stage("experts"); + if (latent_on_) { + D(P_(b.lib, "latent_prompt_pool"), + {mtl::Arg::B(b.YE), mtl::Arg::B(b.pos), mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), + mtl::Arg::B(latent_labels_, labels), mtl::Arg::B(b.lat_pooled), mtl::Arg::V(la)}, + (la.groups * la.tokens * la.width + 255) / 256, 1, 256); + latent_project(du_offset + latent_dense_.d_bytes(), la.width, N, latent_args_.u_row, b.lat_pooled, true); + D(P_(b.lib, "latent_prompt_sum"), {mtl::Arg::B(b.lat_proj), mtl::Arg::B(b.yr), mtl::Arg::V(la)}, + (nT * N + 255) / 256, 1, 256); + } else { + const uint32_t ca[4] = {(uint32_t) nT, N, TOPK, 0}; + D(P_(b.lib, "mm_combine"), {mtl::Arg::B(b.YE), mtl::Arg::B(b.pos), mtl::Arg::B(b.sw), mtl::Arg::B(b.yr), mtl::Arg::V(ca)}, + (nT * N + 255) / 256, 1, 256); + } + end_stage("U"); + gemm(b.mixh, b.Wshg, b.gsh, nT, NFF, N); + gemm(b.mixh, b.Wshu, b.ush, nT, NFF, N); + const uint32_t nh = (uint32_t) nT * NFF; + D(P_(b.lib, "swiglu_h"), {mtl::Arg::B(b.gsh), mtl::Arg::B(b.ush), mtl::Arg::B(b.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); + gemm(b.hh, b.Wshd, b.ys, nT, N, NFF); + gemm(b.mixh, b.Wsg, b.sgv, nT, 1, N); + const uint32_t oa[2] = {(uint32_t) nT, N}; + D(P_(b.lib, "moe_out"), {mtl::Arg::B(b.yr), mtl::Arg::B(b.ys), mtl::Arg::B(b.sgv), mtl::Arg::B(b.bo), mtl::Arg::V(oa)}, + (nT * N + 255) / 256, 1, 256); + end_stage("shared"); + if (l == 0 && ple_on_) { + // layer 0's FFN write applied to R, then the PLE block (as the window: msl/ple.metal over nT rows) + GrArgs a{N, HC, LR, (uint32_t) nT, 1u, eps_}; + D(P_(b.lib, "gr_norm"), {mtl::Arg::B(b.R), mtl::Arg::B(b.bo), mtl::Arg::B(b.inj[cur_inj]), + mtl::Arg::B(W_, T("blk.1.hc_attn_norm.weight").offset), mtl::Arg::B(b.xn), mtl::Arg::V(a)}, + HC, nT, 256); + act16(b.pemb, b.pembh, (size_t) nT * N); + dq(T("blk.1.ple_key.weight"), b.Wpk); + dq(T("blk.1.ple_value.weight"), b.Wpv); + gemm(b.pembh, b.Wpk, b.pkey, nT, 10240, N); + gemm(b.pembh, b.Wpv, b.pval, nT, N, N); + struct PleArgs { uint32_t n_embd, hc; float eps; } pa2{N, HC, strata::kernels::NG_RMS_EPS}; + auto off = [&](const char* n) { return T(n).offset; }; + D(P(lple_, "ple_gnorm", {}, "ple_gnorm"), {mtl::Arg::B(b.pkey), mtl::Arg::B(W_, off("blk.1.ple_norm_key.weight")), + mtl::Arg::B(b.keyn), mtl::Arg::V(pa2)}, HC, nT, 256); + D(P(lple_, "ple_gnorm", {}, "ple_gnorm"), {mtl::Arg::B(b.R), mtl::Arg::B(W_, off("blk.1.ple_norm_query.weight")), + mtl::Arg::B(b.qn), mtl::Arg::V(pa2)}, HC, nT, 256); + D(P(lple_, "ple_gate_norm", {}, "ple_gate_norm"), + {mtl::Arg::B(b.keyn), mtl::Arg::B(b.qn), mtl::Arg::B(b.pval), mtl::Arg::B(W_, off("blk.1.ple_norm_conv.weight")), + mtl::Arg::B(b.gated), mtl::Arg::B(b.normW), mtl::Arg::V(pa2)}, HC, nT, 256); + D(P(lple_, "ple_conv_add", {}, "ple_conv_add"), + {mtl::Arg::B(pleHist_), mtl::Arg::B(b.normW), kw_arg(), mtl::Arg::B(b.gated), + mtl::Arg::B(b.R), mtl::Arg::V(pa2)}, HC * N / 256, nT, 256); + end_stage("PLE"); + } + } + // every row's final residual (the last write applied): the drafter's prompt rows + { + GrArgs a{N, HC, LR, (uint32_t) nT, 1u, eps_}; + D(P_(b.lib, "gr_norm"), {mtl::Arg::B(b.R), mtl::Arg::B(b.bo), mtl::Arg::B(b.inj[cur_inj]), + mtl::Arg::B(W_, T("output_hc_norm.weight").offset), mtl::Arg::B(b.xn), mtl::Arg::V(a)}, HC, nT, 256); + } + if (stage_ends.empty()) { + b.gpu_ms += 1000.0 * cx_->run(ds); + } else { + end_stage("final"); + size_t first = 0; + std::map timings; + for (const auto& [name, last] : stage_ends) { + const std::vector stage(ds.begin() + first, ds.begin() + last); + const double ms = 1000.0 * cx_->run(stage); + timings[name] += ms; + b.gpu_ms += ms; + first = last; + } + std::fprintf(stderr, "strata metal latent prompt profile: %d tokens at %lld", nT, (long long) pos0); + for (const auto& [name, ms] : timings) std::fprintf(stderr, ", %s %.3f ms", name.c_str(), ms); + std::fputc('\n', stderr); + } + hp_ = 1 - hp_; + // the host side of the commit: the PLE history (the last 9 normalized rows) and the n-gram context + if (ple_on_) { + using namespace strata::kernels; + const float* nw = b.normW.as(); + float* hist = pleHist_.as(); + for (int c = 0; c < NG_HC_DIM; ++c) { + float* col = hist + (size_t) c * NG_HIST; + float tmp[NG_HIST * 2]; + int k = 0; + const int keep_old = std::max(0, NG_HIST - nT); + for (int r = NG_HIST - keep_old; r < NG_HIST; ++r) tmp[k++] = col[r]; + for (int t = std::max(0, nT - NG_HIST); t < nT; ++t) tmp[k++] = nw[(size_t) t * NG_HC_DIM + c]; + for (int r = 0; r < NG_HIST; ++r) col[r] = tmp[r]; + } + } + for (int t = 0; t < nT; ++t) { + ss.ple_prev[0] = ss.ple_prev[1]; + ss.ple_prev[1] = (int32_t) tokens[t]; + } + last_T_ = 0; // nothing to commit + return true; +} + +bool Window::prompt(const int64_t* tokens, int64_t n, int64_t pos0, + const std::function& on_chunk, std::string& err) { + static const int C = [] { + const char* e = std::getenv("STRATA_METAL_PROMPT_CHUNK"); + return std::max(64, std::min(8192, e ? std::atoi(e) : 2048)); + }(); + // Reserve the configured latent chunk once. The warm-up sends prompts of + // increasing sizes, and the runtime owns these allocations until teardown. + const int capacity = latent_on_ ? C : (int) std::min(C, std::max(n, 64)); + if (!pc_setup(capacity)) { err = "metal prompt: setup failed"; return false; } + const double t0 = now_ms(); + pb_.gpu_ms = pb_.host_ms = pb_.mtp_ms = 0; + for (int64_t c = 0; c < n; c += C) { + const int nT = (int) std::min(C, n - c); + if (pos0 + c + nT + 1 >= max_cells_) { err = "metal prompt: the prompt runs past the context"; return false; } + if (!pc_chunk(tokens + c, nT, pos0 + c, err)) return false; + if (on_chunk) { + const double tm0 = now_ms(); + mpsrc_ = &pb_.R; // the drafter's prompt rows are this chunk's final residuals + const bool ok = on_chunk(nullptr, nT, pos0 + c, err); + mpsrc_ = nullptr; + pb_.mtp_ms += now_ms() - tm0; + if (!ok) return false; + } + } + if (std::getenv("STRATA_METAL_TRACE") || std::getenv("STRATA_METAL_PROFILE")) + std::fprintf(stderr, "strata metal: prompt %lld tokens in %.1f ms (GPU %.1f ms, host rows %.1f ms, drafter rows %.1f ms) " + "-> %.1f tok/s\n", (long long) n, now_ms() - t0, pb_.gpu_ms, pb_.host_ms, pb_.mtp_ms, + 1000.0 * (double) n / (now_ms() - t0)); + return true; +} diff --git a/edge0-strata/metal/src/engine/setup_host.cpp b/edge0-strata/metal/src/engine/setup_host.cpp new file mode 100644 index 00000000..e4599092 --- /dev/null +++ b/edge0-strata/metal/src/engine/setup_host.cpp @@ -0,0 +1,160 @@ +// metal/src/engine/setup_host.cpp - the host-side state and sizes of the CUDA kernel files that the shared host code +// sets and reads while it starts the engine (path switches, scratch sizes, the RoPE table and its config). Copied +// from gr.cu, shared_expert.cu, native_*.cu, ple.cu, qsa_decode_attn.cu, rope.cu, native_rope.cu, +// s2_expert_grouped.cu; keep in step. On Metal none of these steers a kernel: the decode window runs mac-kern's +// chain (metal/src/engine/window.cpp); they keep the host side's allocations and checks as on the PC. +#include "strata/kernels/gr.hpp" +#include "strata/kernels/native_gdn.hpp" +#include "strata/kernels/native_moe.hpp" +#include "strata/kernels/native_qsa.hpp" +#include "strata/kernels/native_qsa_indexer.hpp" +#include "strata/kernels/native_rope.hpp" +#include "strata/kernels/native_router.hpp" +#include "strata/kernels/ngram.hpp" +#include "strata/kernels/ple.hpp" +#include "strata/kernels/qsa_decode_attn.hpp" +#include "strata/kernels/rope.hpp" +#include "strata/kernels/rope_scaling.hpp" +#include "strata/kernels/s2_expert_grouped.hpp" +#include "strata/kernels/shared_expert.hpp" +#include "strata/kernels/fused_gr.hpp" +#include "strata/prefill/gemm.hpp" + +#include +#include +#include +#include + +namespace strata::kernels { + +namespace { +bool g_fp32_activations = false, g_native_mmvf = false, g_sh_native_bf16 = false, g_ple_native_bf16 = false, + g_ple_native_postops = false; +std::atomic g_moe_combine{false}, g_gdn{false}, g_router{false}, g_qsa{false}, g_qsa_indexer{false}, g_rope{false}; +RopeScaling g_rope_scaling; +} // namespace + +// ---- path switches (each .cu file's own flag) +void gr_set_fp32_activations(bool enabled) { g_fp32_activations = enabled; } +void gr_set_native_mmvf(bool enabled) { g_native_mmvf = enabled; } +bool gr_native_mmvf_enabled() { return g_native_mmvf; } +void shared_expert_set_native_bf16(bool enabled) { g_sh_native_bf16 = enabled; } +bool shared_expert_native_bf16_enabled() { return g_sh_native_bf16; } +void ple_set_native_bf16(bool enabled) { g_ple_native_bf16 = enabled; } +bool ple_native_bf16_enabled() { return g_ple_native_bf16; } +void ple_set_native_postops(bool enabled) { g_ple_native_postops = enabled; } +bool ple_native_postops_enabled() { return g_ple_native_postops; } +void native_moe_combine_set_enabled(bool v) { g_moe_combine.store(v, std::memory_order_relaxed); } +bool native_moe_combine_enabled() { return g_moe_combine.load(std::memory_order_relaxed); } +void native_gdn_set_enabled(bool v) { g_gdn.store(v, std::memory_order_relaxed); } +bool native_gdn_enabled() { return g_gdn.load(std::memory_order_relaxed); } +void native_router_set_enabled(bool v) { g_router.store(v, std::memory_order_relaxed); } +bool native_router_enabled() { return g_router.load(std::memory_order_relaxed); } +void native_qsa_set_enabled(bool v) { g_qsa.store(v, std::memory_order_relaxed); } +bool native_qsa_enabled() { return g_qsa.load(std::memory_order_relaxed); } +void native_qsa_indexer_set_enabled(bool v) { g_qsa_indexer.store(v, std::memory_order_relaxed); } +bool native_qsa_indexer_enabled() { return g_qsa_indexer.load(std::memory_order_relaxed); } +void native_rope_set_enabled(bool v) { g_rope.store(v, std::memory_order_relaxed); } +bool native_rope_enabled() { return g_rope.load(std::memory_order_relaxed); } +void fused_gr_check() {} // the CUDA read variants' self-test: no variants on Metal + +// ---- scratch sizes +uint64_t ple_block_scratch_bytes() { + const size_t f = (size_t) (5 * NG_HC_DIM + NG_N_EMBD + NG_HC) * sizeof(float); + const size_t q = (size_t) (NG_N_EMBD / 32) * 34; + const size_t e = (size_t) NG_N_EMBD * sizeof(uint16_t); + return ((f + 15) & ~(size_t) 15) + ((q + 15) & ~(size_t) 15) + e + 256; +} +uint64_t qsa_decode_attn_scratch_floats(int64_t cap, const QsaShapes& s) { + constexpr int HD = 256, CHUNK = 64; + const int64_t chunks = (cap + CHUNK - 1) / CHUNK; + return (uint64_t) chunks * (uint64_t) s.n_head * (HD + 2) + 64; +} +uint64_t shared_expert_scratch_bytes(int64_t n_ff) { + const uint64_t a = ((uint64_t) n_ff * 4 + 15) & ~15ull; + const uint64_t q0 = ((uint64_t) (n_ff / 32) * 34 + 15) & ~15ull; + const uint64_t qk = ((uint64_t) (n_ff / 256) * 292 + 15) & ~15ull; + return a * 2 + q0 + qk + 32; +} +uint64_t moe_hit_grouped_scratch_bytes(int64_t n_hits, int64_t n_embd, int64_t n_ff) { + if (n_hits <= 0) return 0; + const uint64_t gu = (uint64_t) n_hits * (uint64_t) (2 * n_ff) * 4; + const uint64_t q8 = (uint64_t) n_hits * (uint64_t) (n_ff / 32) * 34; + const uint64_t hs = (uint64_t) n_hits * (uint64_t) (n_ff / 32) * 4; + const uint64_t xh = (uint64_t) (n_embd / 32) * 4; + return ((gu + 15) & ~15ull) + ((q8 + 15) & ~15ull) + 2 * ((hs + 15) & ~15ull) + ((xh + 15) & ~15ull); +} +size_t gr_workspace_init(const GrShapes& s, void* base, GrWorkspace& out) { + const size_t hc_dim = (size_t) s.hc * (size_t) s.n_embd; + const size_t sz[5] = {hc_dim * sizeof(float), hc_dim * sizeof(uint16_t), (size_t) s.hc_lr * sizeof(uint16_t), + hc_dim * sizeof(float), (size_t) s.hc_lr * sizeof(float)}; + size_t al[5], bytes = 0; + for (int k = 0; k < 5; ++k) { + al[k] = (sz[k] + 15) & ~(size_t) 15; + bytes += al[k]; + } + out.bytes = bytes; + if (base != nullptr) { + unsigned char* p = (unsigned char*) base; + void* ptr[5]; + for (int k = 0; k < 5; ++k) { + ptr[k] = p; + p += al[k]; + } + out.xn = (float*) ptr[0]; + out.xq = (uint16_t*) ptr[1]; + out.lq = (uint16_t*) ptr[2]; + out.gated = (float*) ptr[3]; + out.lo = (float*) ptr[4]; + } + return bytes; +} + +// ---- RoPE (rope.cu, native_rope.cu): the config and the host-built angle table +void rope_scaling_set(const RopeScaling& scaling) { g_rope_scaling = scaling; } +const RopeScaling& rope_scaling() { return g_rope_scaling; } +void build_rope_table(int n_rot, double theta, int max_pos, float* cos_tab, float* sin_tab) { + const int half = n_rot / 2; + for (int p = 0; p < max_pos; ++p) + for (int i = 0; i < half; ++i) { + const double inv = std::pow(theta, -2.0 * (double) i / (double) n_rot); + const double ang = (double) p * inv; + cos_tab[(size_t) p * half + i] = (float) std::cos(ang); + sin_tab[(size_t) p * half + i] = (float) std::sin(ang); + } +} +void build_rope_table(int n_rot, const RopeScaling& sc, int max_pos, float* cos_tab, float* sin_tab) { + if (sc.type == RopeScalingType::None) { + build_rope_table(n_rot, sc.freq_base, max_pos, cos_tab, sin_tab); + return; + } + const int half = n_rot / 2; + const double fs = sc.freq_scale(); + const double ms = sc.mscale(); + double cd[2]; + sc.corr_dims(n_rot, cd); + const bool correct = sc.ext_factor != 0; + for (int p = 0; p < max_pos; ++p) + for (int i = 0; i < half; ++i) { + const double inv = std::pow(sc.freq_base, -2.0 * (double) i / (double) n_rot); + const double extrap = (double) p * inv; + const double interp = fs * extrap; + double ang = interp; + if (correct) { + const double ramp = (double) rope_yarn_ramp((float) cd[0], (float) cd[1], i) * sc.ext_factor; + ang = interp * (1.0 - ramp) + extrap * ramp; + } + cos_tab[(size_t) p * half + i] = (float) (std::cos(ang) * ms); + sin_tab[(size_t) p * half + i] = (float) (std::sin(ang) * ms); + } +} +void rope_table_set(const float*, const float*, int, const RopeScaling&) {} // the window keeps its own table +void rope_table_release(const float*) {} +void native_q6_k_unpack(const void*) {} // nothing packed on Metal + +} // namespace strata::kernels + +namespace strata::prefill { +Gemm::~Gemm() {} // nothing created on Metal (no prompt GEMMs yet) +bool prompt_f16() { return false; } // the CUDA/HIP prompt path's FP16 image mode: not used on Metal +} // namespace strata::prefill diff --git a/edge0-strata/metal/src/engine/stream.cpp b/edge0-strata/metal/src/engine/stream.cpp new file mode 100644 index 00000000..4d90b985 --- /dev/null +++ b/edge0-strata/metal/src/engine/stream.cpp @@ -0,0 +1,824 @@ +// metal/src/engine/stream.cpp - see stream.h. +#include "stream.h" + +#include "ref_quant.h" +#include "arena_budget.h" + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace strata::metal_engine { + +namespace { +constexpr int NL = 48, NE = 512; +constexpr uint64_t PAGE = 16384; +double now_ms() { return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); } +int env_int(const char* n, int d) { + const char* e = std::getenv(n); + return e && *e ? std::atoi(e) : d; +} +} // namespace + +bool ExpertStream::wanted(const gguf::File& f, const mtl::Context& cx) { + const char* e = std::getenv("STRATA_METAL_STREAM"); + if (e && std::string(e) == "1") return true; + if (e && std::string(e) == "0") return false; + // auto: stream when the mapped weights (experts and dense) and ~12 GiB beside them (the prompt buffers, KV, the + // draft layer, macOS) do not fit the resident footprint cap: the GPU working set, or on a big Mac the RAM less the + // 30% macOS and its file cache keep (arena_budget.h: footprint_cap). M5 Max (115.4 GB working set, 96.2 GB cap): + // GSQ packs, IQ3_S and UD-Q4_K_XL stay resident, UD-Q5/Q6_K_XL stream; M4 32 GB (26.8 GB): every pack streams. + uint64_t all = 0; + for (const auto& kv : f.tensors) + if (kv.first != "per_layer_token_embd.weight") all += kv.second.bytes; + uint64_t phys = 0; + size_t len = sizeof phys; + sysctlbyname("hw.memsize", &phys, &len, nullptr, 0); + return !resident_fits((double) phys, (double) cx.recommended_working_set(), (double) all, 12.0 * (1ull << 30)); +} + +ExpertStream::~ExpertStream() { + { + std::lock_guard g(mu_); + stop_ = true; + } + cv_.notify_all(); + for (auto& t : io_) t.join(); + if (fd_ >= 0) ::close(fd_); +} + +bool ExpertStream::open(const gguf::File& f, std::string& err) { + fd_ = ::open(f.path.c_str(), O_RDONLY); + if (fd_ < 0) { err = "metal stream: cannot open " + f.path; return false; } + fcntl(fd_, F_NOCACHE, 1); + fcntl(fd_, F_RDAHEAD, 0); + lay_.resize(NL); + for (int l = 0; l < NL; ++l) { + const std::string b = "blk." + std::to_string(l) + "."; + const gguf::Tensor* g = f.find(b + "ffn_gate_exps.weight"); + const gguf::Tensor* u = f.find(b + "ffn_up_exps.weight"); + const gguf::Tensor* d = f.find(b + "ffn_down_exps.weight"); + if (!g || !u || !d) { err = "metal stream: layer " + std::to_string(l) + " has no routed experts"; return false; } + lay_[(size_t) l] = {g->offset, g->bytes / NE, u->offset, u->bytes / NE, d->offset, d->bytes / NE, -1}; + } + // the expert-contiguous copy beside the GGUF, when it is this GGUF's (sizes in its header) and its first and + // last experts match the GGUF's bytes + if (env_int("STRATA_METAL_CONTIG", 1) != 0) { + const std::string cp = f.path + ".experts"; + const int cfd = ::open(cp.c_str(), O_RDONLY); + bool ok = cfd >= 0; + std::vector off(NL); + if (ok) { + uint8_t h[64]; + std::vector t((size_t) NL * 4); + struct stat st {}; + ::fstat(fd_, &st); + ok = ::pread(cfd, h, 64, 0) == 64 && std::memcmp(h, "STRX", 4) == 0 && + ::pread(cfd, t.data(), t.size() * 8, 64) == (ssize_t) (t.size() * 8); + uint32_t ver = 0, nl = 0, ne = 0; + uint64_t src = 0; + std::memcpy(&ver, h + 4, 4); std::memcpy(&nl, h + 8, 4); std::memcpy(&ne, h + 12, 4); std::memcpy(&src, h + 16, 8); + ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) NE && src == (uint64_t) st.st_size; + for (int l = 0; ok && l < NL; ++l) { + const Lay& x = lay_[(size_t) l]; + ok = t[(size_t) l * 4 + 1] == x.gs && t[(size_t) l * 4 + 2] == x.us && t[(size_t) l * 4 + 3] == x.ds; + off[(size_t) l] = t[(size_t) l * 4]; + } + if (ok) { // spot check: layer 0's first and layer 47's last expert, every slice + for (const int l : {0, NL - 1}) { + const Lay& x = lay_[(size_t) l]; + const int e = l == 0 ? 0 : NE - 1; + std::vector a(x.gs + x.us + x.ds), b(a.size()); + ok = ok && ::pread(cfd, a.data(), a.size(), (off_t) (off[(size_t) l] + (uint64_t) e * a.size())) == (ssize_t) a.size() && + ::pread(fd_, b.data(), x.gs, (off_t) (x.g + (uint64_t) e * x.gs)) == (ssize_t) x.gs && + ::pread(fd_, b.data() + x.gs, x.us, (off_t) (x.u + (uint64_t) e * x.us)) == (ssize_t) x.us && + ::pread(fd_, b.data() + x.gs + x.us, x.ds, (off_t) (x.d + (uint64_t) e * x.ds)) == (ssize_t) x.ds && + a == b; + } + } + } + if (ok) { + fcntl(cfd, F_NOCACHE, 1); + fcntl(cfd, F_RDAHEAD, 0); + cfd_ = cfd; + coff_ = off; + std::fprintf(stderr, "strata metal stream: experts read from %s (one request per expert)\n", cp.c_str()); + } else if (cfd >= 0) { + ::close(cfd); + std::fprintf(stderr, "strata metal stream: %s is not this GGUF's (sizes or bytes differ): the GGUF is read\n", + cp.c_str()); + } + } + // size classes: one LRU pool per class. The layers' slot sizes (9 in the IQ3_XXS mix) are merged into a few + // classes, each as big as its largest member, so that no member wastes more than STRATA_METAL_CLASS_WASTE (default + // 0.15) of its slot: every class must hold a whole layer's experts for a prompt chunk, so fewer classes = less + // memory tied up in those minimums. + { + const char* ew = std::getenv("STRATA_METAL_CLASS_WASTE"); + const double waste = ew ? std::atof(ew) : 0.15; + std::vector pages_of(NL); + std::vector sizes; + for (int l = 0; l < NL; ++l) { + const Lay& x = lay_[(size_t) l]; + pages_of[(size_t) l] = (x.gs + x.us + x.ds + PAGE - 1) / PAGE; + sizes.push_back(pages_of[(size_t) l]); + } + std::sort(sizes.begin(), sizes.end()); + sizes.erase(std::unique(sizes.begin(), sizes.end()), sizes.end()); + std::vector tops; // each class's slot pages, largest first + for (auto it = sizes.rbegin(); it != sizes.rend(); ++it) + if (tops.empty() || (double) *it < (1.0 - waste) * (double) tops.back()) tops.push_back(*it); + for (uint64_t t : tops) { Cls c; c.pages = t; cls_.push_back(c); } + for (int l = 0; l < NL; ++l) { + int best = 0; // the smallest class that holds the layer's slot + for (int c = 0; c < (int) cls_.size(); ++c) + if (cls_[(size_t) c].pages >= pages_of[(size_t) l]) best = c; + lay_[(size_t) l].cls = best; + } + } + nio_ = std::max(1, env_int("STRATA_METAL_IO_THREADS", 6)); + split_ = std::max(1, std::min(8, env_int("STRATA_METAL_READ_SPLIT", 1))); + // The read-ahead's depth and where its experts wait in the LRU, measured on the M4 (notes-mac-ssd.md, 3 interleaved + // pairs each): Q2_0 (1.38 MB experts) top 10 at the LRU front 12.2 tok/s vs top 6 at the tail 11.3; the IQ3_XXS + // mix (1.75 MB, the SSD is the limit) top 6 at the tail 8.06 vs top 10 at the front 7.43. + uint64_t eb = 0; + for (const auto& x : lay_) eb += x.gs + x.us + x.ds; + const bool big = (double) eb / NL > 1.5e6; + pre_k_ = env_int("STRATA_METAL_PREFETCH", big ? 6 : 10); + pre_k2_ = env_int("STRATA_METAL_PREFETCH2", 0); + prefetch_ = pre_k_ > 0; + pre_mru_ = env_int("STRATA_METAL_PREFETCH_MRU", big ? 0 : 1) != 0; + // timing only (WRONG answers): every missing expert is pointed at slot 0, nothing is read + fake_ = env_int("STRATA_METAL_STREAM_FAKE", 0) != 0; + return true; +} + +bool ExpertStream::read(void* dst, uint64_t len, uint64_t off) const { + uint8_t* p = (uint8_t*) dst; + while (len > 0) { + const ssize_t r = ::pread(fd_, p, (size_t) std::min(len, 64ull << 20), (off_t) off); + if (r <= 0) return false; + p += r; + off += (uint64_t) r; + len -= (uint64_t) r; + } + return true; +} + +bool ExpertStream::load_dense(gguf::File& f, mtl::Context& cx, const std::vector& skip, mtl::Buffer& out, + std::string& err) { + const double t0 = now_ms(); + std::vector ts; + uint64_t tot = 0; + for (auto& kv : f.tensors) { + gguf::Tensor& t = kv.second; + if (t.name.find("_exps") != std::string::npos) continue; + // the n-gram table is read by the host from the SSD, a few rows at a time (src/ngram/ple_reader.cpp). It is its own shard in + // the ISTA / Unsloth files; in agentionai's Gyro files it is the first 28.8 GB of the one file, and must not be read in here + if (t.name == "per_layer_token_embd.weight") continue; + if (std::find(skip.begin(), skip.end(), t.name) != skip.end()) { t.offset = 0; continue; } + ts.push_back(&t); + tot += (t.bytes + 255) / 256 * 256; + } + out = cx.alloc((tot + PAGE - 1) / PAGE * PAGE); + std::vector src(ts.size()), dst(ts.size()); + uint64_t o = 0; + for (size_t i = 0; i < ts.size(); ++i) { + src[i] = ts[i]->offset; + dst[i] = o; + o += (ts[i]->bytes + 255) / 256 * 256; + } + std::atomic next{0}; + std::atomic ok{true}; + std::vector th; + for (int k = 0; k < 4; ++k) + th.emplace_back([&] { + for (;;) { + const size_t i = next++; + if (i >= ts.size()) break; + if (!read(out.as() + dst[i], ts[i]->bytes, src[i])) ok = false; + } + }); + for (auto& x : th) x.join(); + if (!ok) { err = "metal stream: reading the dense tensors failed"; return false; } + for (size_t i = 0; i < ts.size(); ++i) ts[i]->offset = dst[i]; + std::fprintf(stderr, "strata metal stream: %zu dense tensors read into memory (%.2f GB, %.1f s; %zu left out)\n", ts.size(), + tot / 1e9, (now_ms() - t0) / 1e3, skip.size()); + return true; +} + +int ExpertStream::cls_of(int s) const { // the class ranges are contiguous + int ci = 0; + while (!(s >= cls_[(size_t) ci].base && s < cls_[(size_t) ci].base + cls_[(size_t) ci].n)) ++ci; + return ci; +} + +void ExpertStream::lru_unlink(int s) { + Slot& x = slots_[(size_t) s]; + Cls& k = cls_[(size_t) cls_of(s)]; + if (x.prev >= 0) slots_[(size_t) x.prev].next = x.next; else k.head = x.next; + if (x.next >= 0) slots_[(size_t) x.next].prev = x.prev; else k.tail = x.prev; + x.prev = x.next = -1; +} + +void ExpertStream::lru_back(int s) { + lru_unlink(s); + Cls& k = cls_[(size_t) cls_of(s)]; + Slot& x = slots_[(size_t) s]; + x.prev = k.tail; + x.next = -1; + if (k.tail >= 0) slots_[(size_t) k.tail].next = s; + k.tail = s; + if (k.head < 0) k.head = s; +} + +void ExpertStream::lru_front(int s) { + lru_unlink(s); + Cls& k = cls_[(size_t) cls_of(s)]; + Slot& x = slots_[(size_t) s]; + x.next = k.head; + x.prev = -1; + if (k.head >= 0) slots_[(size_t) k.head].prev = s; + k.head = s; + if (k.tail < 0) k.tail = s; +} + +int ExpertStream::victim(int ci) { + Cls& k = cls_[(size_t) ci]; + for (int s = k.tail; s >= 0; s = slots_[(size_t) s].prev) + if (slots_[(size_t) s].state.load() == 0 && slots_[(size_t) s].pin != stamp_) { + Slot& x = slots_[(size_t) s]; + if (x.layer >= 0) { + const int key = x.layer * NE + x.e; + where_[(size_t) key] = -1; + smap_.as()[key] = 0xFFFFFFFFu; + } + x.layer = x.e = -1; + return s; + } + return -1; +} + +// One piece of an expert: slice part / split_ (gate, up, down), piece part % split_ of it. The pieces of one expert +// go to several I/O threads, so a single miss is read with several requests in flight (its latency, not only the +// SSD's throughput, is what a layer waits for). +void ExpertStream::read_part(int s, int l, int e, int part) { + const Lay& x = lay_[(size_t) l]; + if (cfd_ >= 0) { // the contiguous copy: the whole expert, piece part of split_ (the slot keeps gate | up | down) + const uint64_t len = x.gs + x.us + x.ds, src = coff_[(size_t) l] + (uint64_t) e * len; + const uint64_t a = len * (uint64_t) part / (uint64_t) split_ / 4096 * 4096; + const uint64_t b = part + 1 == split_ ? len : len * (uint64_t) (part + 1) / (uint64_t) split_ / 4096 * 4096; + uint8_t* p = host_[(size_t) s] + a; + uint64_t n = b - a, o = src + a; + while (n > 0) { + const ssize_t r = ::pread(cfd_, p, (size_t) n, (off_t) o); + if (r <= 0) { + std::fprintf(stderr, "strata metal stream: reading expert %d of layer %d failed\n", e, l); + std::exit(6); + } + p += r; o += (uint64_t) r; n -= (uint64_t) r; + } + bytes_read_ += b - a; + return; + } + const int k = part / split_, p = part % split_; + const uint64_t len = k == 0 ? x.gs : (k == 1 ? x.us : x.ds); + const uint64_t src = (k == 0 ? x.g : (k == 1 ? x.u : x.d)) + (uint64_t) e * len; + const uint64_t dst = k == 0 ? 0 : (k == 1 ? x.gs : x.gs + x.us); + const uint64_t a = len * (uint64_t) p / (uint64_t) split_ / 4096 * 4096; + const uint64_t b = p + 1 == split_ ? len : len * (uint64_t) (p + 1) / (uint64_t) split_ / 4096 * 4096; + if (b > a && !read(host_[(size_t) s] + dst + a, b - a, src + a)) { + std::fprintf(stderr, "strata metal stream: reading expert %d of layer %d failed\n", e, l); + std::exit(6); + } + bytes_read_ += b - a; +} + +void ExpertStream::io_main() { + for (;;) { + Job j; + { + std::unique_lock g(mu_); + cv_.wait(g, [&] { return stop_ || !urgent_q_.empty() || !pre_q_.empty(); }); + if (stop_) return; + if (!urgent_q_.empty()) { j = urgent_q_.front(); urgent_q_.pop_front(); } + else { j = pre_q_.front(); pre_q_.pop_front(); } + slots_[(size_t) j.slot].state = 2; + } + read_part(j.slot, j.layer, j.e, j.part); + bool done = false; + { + std::lock_guard g(mu_); + Slot& x = slots_[(size_t) j.slot]; + if (--x.left == 0) { + smap_.as()[j.layer * NE + j.e] = code_[(size_t) j.slot]; + x.state = 0; + --outstanding_; + ++reads_; + done = true; + } + } + if (done) done_cv_.notify_all(); + } +} + +void ExpertStream::submit(int s, int l, int e, bool urgent) { + Slot& x = slots_[(size_t) s]; + x.layer = (int16_t) l; + x.e = (int16_t) e; + x.state = 1; + x.pre = !urgent; + const int parts = cfd_ >= 0 ? split_ : 3 * split_; + x.left = parts; + ++outstanding_; + where_[(size_t) (l * NE + e)] = s; + for (int p = 0; p < parts; ++p) (urgent ? urgent_q_ : pre_q_).push_back({s, l, e, p}); +} + +// a queued read-ahead nobody wants: its pieces dropped, the slot freed (only before any piece has started) +void ExpertStream::cancel_slot(int s) { + Slot& x = slots_[(size_t) s]; + for (auto it = pre_q_.begin(); it != pre_q_.end();) it = it->slot == s ? pre_q_.erase(it) : it + 1; + where_[(size_t) (x.layer * NE + x.e)] = -1; + x.layer = x.e = -1; + x.state = 0; + x.pre = false; + x.left = 0; + --outstanding_; + ++pre_cancel_; + lru_back(s); +} + +// a read-ahead that is wanted now: its queued pieces go first +void ExpertStream::promote_slot(int s) { + std::vector mine; + for (auto it = pre_q_.begin(); it != pre_q_.end();) + if (it->slot == s) { mine.push_back(*it); it = pre_q_.erase(it); } + else ++it; + for (auto it = mine.rbegin(); it != mine.rend(); ++it) urgent_q_.push_front(*it); +} + +bool ExpertStream::make_arena(mtl::Context& cx, uint64_t budget, const std::string& profile, int min_slots, std::string& err, + const std::function& recheck, const std::function& grow) { + const double t0 = now_ms(); + shrink_to_ = 0; + for (auto& k : cls_) k.n = 0; // a second try (after a shrink) lays the arena out again + // the ranking: the profile's pairs, then every pair it misses, layer-interleaved + std::vector> rank; + std::vector seen((size_t) NL * NE, 0); + if (FILE* f = std::fopen(profile.c_str(), "rb")) { + char magic[4]; + uint32_t h[5]; + if (std::fread(magic, 1, 4, f) == 4 && std::memcmp(magic, "STRP", 4) == 0 && std::fread(h, 4, 5, f) == 5 && + h[1] == (uint32_t) NL && h[2] == (uint32_t) NE) { + for (uint32_t i = 0; i < h[4]; ++i) { + uint16_t p[2]; + if (std::fread(p, 2, 2, f) != 2) break; + if (p[0] < NL && p[1] < NE && !seen[(size_t) p[0] * NE + p[1]]) { + seen[(size_t) p[0] * NE + p[1]] = 1; + rank.push_back({p[0], p[1]}); + } + } + } + std::fclose(f); + } else { + std::fprintf(stderr, "strata metal stream: no expert profile at %s: the arena starts layer-interleaved\n", profile.c_str()); + } + const size_t from_profile = rank.size(); + for (int e = 0; e < NE; ++e) + for (int l = 0; l < NL; ++l) + if (!seen[(size_t) l * NE + e]) rank.push_back({l, e}); + // the pool of each size class: at least min_slots (a prompt chunk's layer needs up to 512 at once), then the + // top-ranked pairs that fit in the budget + uint64_t used = 0; + for (const auto& k : cls_) used += (uint64_t) min_slots * k.pages * PAGE; + std::vector> fill; + for (const auto& p : rank) { + Cls& k = cls_[(size_t) lay_[(size_t) p.first].cls]; + const uint64_t cost = k.n < min_slots ? 0 : k.pages * PAGE; + if (used + cost > budget) break; + used += cost; + ++k.n; + fill.push_back(p); + } + int base = 0; + for (auto& k : cls_) { + k.n = std::max(k.n, min_slots); + k.base = base; + base += k.n; + } + // the slots placed in chunks of at most maxBufferLength (and 8 GB), up to 4 (the kernels' buffers 0, 7, 8, 9) + // (8 GB each unless the arena needs bigger ones to fit 4: the M4 refused one 20.6 GB buffer, the M5 Max takes an + // 89 GB arena for UD-Q5_K_XL in 4 of ~23 GB; STRATA_METAL_ARENA_CHUNK_GB sets it) + uint64_t total = 0, biggest = 0; + for (const auto& k : cls_) { + total += (uint64_t) k.n * k.pages * PAGE; + biggest = std::max(biggest, k.pages * PAGE); + } + uint64_t want = std::max(8ull << 30, total / 4 + 4 * biggest); + if (const char* e = std::getenv("STRATA_METAL_ARENA_CHUNK_GB")) want = (uint64_t) (std::atof(e) * (1ull << 30)); + const uint64_t cmax = std::min(cx.max_buffer_length(), want) / PAGE * PAGE; + code_.assign((size_t) base, 0); + host_.assign((size_t) base, nullptr); + std::vector csize; + std::vector> at((size_t) base); + uint64_t cur = 0; + for (const auto& k : cls_) + for (int s = k.base; s < k.base + k.n; ++s) { + if (csize.empty() || cur + k.pages * PAGE > cmax) { + if (!csize.empty()) csize.back() = cur; + csize.push_back(0); + cur = 0; + } + at[(size_t) s] = {(int) csize.size() - 1, cur / PAGE}; + cur += k.pages * PAGE; + } + if (!csize.empty()) csize.back() = cur; + if (csize.size() > 4) { err = "metal stream: the arena needs more than 4 buffers (lower STRATA_METAL_ARENA_GB)"; return false; } + if (recheck) { // the memory macOS can give, asked again just before the arena is committed + uint64_t total = 0; + for (uint64_t b : csize) total += b; + const uint64_t allowed = recheck(); + if (allowed + (64ull << 20) < total) { + shrink_to_ = allowed; + err = "metal stream: less memory than planned is free now: the arena shrinks"; + return false; + } + } + for (uint64_t b : csize) { // anonymous memory wrapped no-copy: a prompt's buffers can alias a region of it + void* m = mmap(nullptr, (size_t) b, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0); + if (m == MAP_FAILED) { err = "metal stream: cannot map the arena"; return false; } + chunks_.push_back(cx.wrap(m, (size_t) b)); + arena_bytes_ += b; + } + for (int s = 0; s < base; ++s) { + code_[(size_t) s] = ((uint32_t) at[(size_t) s].first << 28) | (uint32_t) at[(size_t) s].second; + host_[(size_t) s] = chunks_[(size_t) at[(size_t) s].first].as() + at[(size_t) s].second * PAGE; + } + smap_ = cx.alloc((size_t) NL * NE * 4); + std::memset(smap_.ptr, 0xFF, (size_t) NL * NE * 4); + slots_ = std::vector((size_t) base); + retired_.assign((size_t) base, 0); + where_.assign((size_t) NL * NE, -1); + for (auto& k : cls_) { + for (int s = k.base; s < k.base + k.n; ++s) { + slots_[(size_t) s].prev = s == k.base ? -1 : s - 1; + slots_[(size_t) s].next = s + 1 == k.base + k.n ? -1 : s + 1; + } + k.head = k.n ? k.base : -1; + k.tail = k.n ? k.base + k.n - 1 : -1; + } + for (int i = 0; i < nio_; ++i) io_.emplace_back([this] { io_main(); }); + // the fill: the ranked pairs in steps (best first; each step in file order) on the I/O threads; before each step + // grow() asks whether macOS still has the memory for it (arena_grow); the best-ranked ends at the LRU front + uint64_t step = 1536ull << 20; + if (const char* e = std::getenv("STRATA_METAL_ARENA_STEP_GB")) step = std::max(64ull << 20, (uint64_t) (std::atof(e) * 1e9)); + std::vector next_free(cls_.size()); + for (size_t c = 0; c < cls_.size(); ++c) next_free[c] = cls_[c].base; + size_t done = 0; + bool stopped = false; + uint64_t filled = 0; + while (done < fill.size()) { + size_t j = done; + uint64_t sb = 0; + while (j < fill.size() && (sb < step || j == done)) sb += cls_[(size_t) lay_[(size_t) fill[j++].first].cls].pages * PAGE; + if (grow && !grow(sb)) { stopped = true; break; } + std::vector> order(fill.begin() + (long) done, fill.begin() + (long) j); + std::sort(order.begin(), order.end(), [&](const auto& a, const auto& b) { + return lay_[(size_t) a.first].g + (uint64_t) a.second * lay_[(size_t) a.first].gs < + lay_[(size_t) b.first].g + (uint64_t) b.second * lay_[(size_t) b.first].gs; + }); + { + std::lock_guard g(mu_); + for (const auto& p : order) { + const int c = lay_[(size_t) p.first].cls; + submit(next_free[(size_t) c]++, p.first, p.second, true); + } + } + cv_.notify_all(); + { + std::unique_lock g(mu_); + done_cv_.wait(g, [&] { return outstanding_ == 0; }); + } + filled += sb; + done = j; + } + uint64_t retired_bytes = 0; + size_t nretired = 0; + { + std::unique_lock g(mu_); + if (stopped) // the slots not filled stay unused (their pages never touched), all but min_slots of each class + for (size_t c = 0; c < cls_.size(); ++c) { + const Cls& k = cls_[c]; + for (int s = std::max(next_free[c], k.base + std::min(min_slots, k.n)); s < k.base + k.n; ++s) { + lru_unlink(s); + retired_[(size_t) s] = 1; + retired_bytes += k.pages * PAGE; + ++nretired; + } + } + // the LRU order = the rank order (best at the front) + for (size_t i = done; i-- > 0;) lru_front(where_[(size_t) (fill[i].first * NE + fill[i].second)]); + } + if (stopped) + std::fprintf(stderr, "strata metal stream: the arena's fill stopped after %.2f GB: macOS's free memory reached what it " + "keeps; %.2f GB of the arena (%zu slots) stay unused\n", filled / 1e9, retired_bytes / 1e9, nretired); + arena_live_ = arena_bytes_ - retired_bytes; + fill_s_ = (now_ms() - t0) / 1e3; + fill_bytes_ = bytes_read_.load(); + uint64_t pairs = 0; + for (const auto& k : cls_) pairs += (uint64_t) k.n; + std::fprintf(stderr, "strata metal stream: expert arena %.2f GB, %llu slots in %zu size classes, %zu experts read from the " + "profile's top (%zu ranked) in %.1f s (%.2f GB/s); %d I/O threads, read-ahead top %d; %zu buffers\n", + (arena_bytes_ - retired_bytes) / 1e9, (unsigned long long) (pairs - nretired), cls_.size(), done, from_profile, fill_s_, + fill_bytes_ / 1e9 / std::max(1e-9, fill_s_), nio_, prefetch_ ? pre_k_ : 0, chunks_.size()); + // Wire the filled arena (mlock). Unwired, the arena is ordinary anonymous memory whenever the GPU is not running a command buffer + // (the GPU's own wiring ends a few seconds after the last one; a Metal residency set did not hold it either, measured), and + // macOS compresses it when anything else wants memory: a 20 GB file written beside the idle engine took tens of GB of the arena + // into the compressor in seconds and froze the 128 GB M5 Max (watchdog panic, 2026-10-13). Wired, a file write can only + // shrink the file cache. STRATA_METAL_WIRE_ARENA=0/1; the default is on for a Mac of 64 GB or more (the M4 32 GB was measured + // unwired). Not when the fill stopped early for lack of memory (the unused slots would be wired too). + { + uint64_t phys = 0; + size_t plen = sizeof phys; + sysctlbyname("hw.memsize", &phys, &plen, nullptr, 0); + const char* we = std::getenv("STRATA_METAL_WIRE_ARENA"); + const bool want = we && *we ? std::atoi(we) != 0 : (double) phys >= kBigMac; + if (want && !stopped) { + const double tp = now_ms(); + uint64_t pinned = 0; + int failed = 0; + for (const auto& c : chunks_) { + if (mlock(c.ptr, c.bytes) == 0) pinned += c.bytes; + else ++failed; + } + std::fprintf(stderr, "strata metal stream: %.1f GB of the arena wired (mlock) in %.0f ms%s; macOS cannot compress it while " + "the engine idles (STRATA_METAL_WIRE_ARENA=0 turns it off)\n", pinned / 1e9, now_ms() - tp, + failed ? "; SOME CHUNKS COULD NOT BE WIRED" : ""); + } + } + (void) err; + return true; +} + +// the experts a predicted routing names, most likely first (by the best logit over the window's tokens) +void ExpertStream::ahead(int ln, int T, const float* pred, int K) { + K = std::min(K, 32); + std::pair pw[8 * 32]; + int np = 0; + for (int t = 0; t < T; ++t) { + const float* lg = pred + (size_t) t * NE; + int top[32]; + int nt = 0; + for (int e = 0; e < NE; ++e) { + if (nt < K) { top[nt++] = e; continue; } + int m = 0; + for (int j = 1; j < K; ++j) if (lg[top[j]] < lg[top[m]]) m = j; + if (lg[e] > lg[top[m]]) top[m] = e; + } + for (int j = 0; j < nt; ++j) { + int q = 0; + while (q < np && pw[q].second != top[j]) ++q; + if (q == np) pw[np++] = {lg[top[j]], top[j]}; + else pw[q].first = std::max(pw[q].first, lg[top[j]]); + } + } + std::sort(pw, pw + np, [](const auto& a, const auto& b) { return a.first > b.first; }); + bool any = false; + { + std::lock_guard g(mu_); + for (int j = 0; j < np; ++j) { + const int key = ln * NE + pw[j].second; + const int s0 = where_[(size_t) key]; + if (s0 >= 0) { lru_front(s0); continue; } + const int s = victim(lay_[(size_t) ln].cls); + if (s < 0) break; + submit(s, ln, pw[j].second, false); + slots_[(size_t) s].pin = stamp_; + if (pre_mru_) lru_front(s); + else lru_back(s); + ++pre_issued_; + any = true; + } + } + if (any) cv_.notify_all(); +} + +void ExpertStream::at_layer(int l, int T, int topk, const int32_t* ids, const float* pred, const float* pred2) { + const double t0 = now_ms(); + const bool prompt = T > 8; // a prompt chunk (the GEMM path): nearly every expert of the layer is used + if (l == 0) { if (prompt) ++p_chunks_; else ++windows_; } + ++stamp_; // the slots this layer uses are pinned until the next handshake + bool used[NE] = {}; + int want[NE]; + int nw = 0; + for (int i = 0; i < T * topk; ++i) { + const int e = ids[i]; + if (e < 0 || e >= NE || used[e]) continue; + used[e] = true; + want[nw++] = e; + } + if (fake_) { + for (int j = 0; j < nw; ++j) + if (where_[(size_t) (l * NE + want[j])] < 0) smap_.as()[l * NE + want[j]] = code_[0]; + host_ms_ += now_ms() - t0; + return; + } + std::vector wait; + { + std::lock_guard g(mu_); + // read-aheads of this layer (or older) that were not started and are not wanted: dropped (the slot goes free) + { + std::vector drop; + for (const Job& j : pre_q_) { + if (j.layer > l || slots_[(size_t) j.slot].state.load() != 1) continue; + if (!(j.layer == l && used[j.e]) && std::find(drop.begin(), drop.end(), j.slot) == drop.end()) drop.push_back(j.slot); + } + for (int s : drop) cancel_slot(s); + } + for (int j = 0; j < nw; ++j) // pass 1: pin what is resident (or on its way), so no miss below evicts it + if (where_[(size_t) (l * NE + want[j])] >= 0) slots_[(size_t) where_[(size_t) (l * NE + want[j])]].pin = stamp_; + for (int j = 0; j < nw; ++j) { + const int e = want[j], key = l * NE + e; + uint64_t& lk = prompt ? p_lookups_ : lookups_; + ++lk; + int s = where_[(size_t) key]; + if (s >= 0) { + Slot& x = slots_[(size_t) s]; + x.pin = stamp_; + const int st = x.state.load(); + if (st == 0) { if (x.pre) ++(prompt ? p_pre_ : pre_hits_); else ++(prompt ? p_hits_ : hits_); } + else { + ++(prompt ? p_pre_ : pre_hits_); + if (!prompt) ++late_; + promote_slot(s); // a queued read-ahead becomes urgent + wait.push_back(s); + } + if (!(prompt && x.pre)) lru_front(s); // a prompt's streamed experts stay at the LRU tail + x.pre = false; + continue; + } + ++(prompt ? p_misses_ : misses_); + s = victim(lay_[(size_t) l].cls); + if (s < 0) { std::fprintf(stderr, "strata metal stream: no free slot (layer %d)\n", l); std::exit(6); } + submit(s, l, e, true); + slots_[(size_t) s].pin = stamp_; + if (!prompt) lru_front(s); + wait.push_back(s); + } + } + if (!wait.empty()) { + cv_.notify_all(); + const double tw = now_ms(); + std::unique_lock g(mu_); + done_cv_.wait(g, [&] { + for (int s : wait) if (slots_[(size_t) s].state.load() != 0) return false; + return true; + }); + (prompt ? p_wait_ms_ : wait_ms_) += now_ms() - tw; + } + if (prompt) { + // a chunk uses nearly all of the next layer's experts: every one not in the arena is streamed in behind this + // layer's compute, into the LRU tail (the decode-hot experts stay) + if (l + 1 < NL) { + bool any = false; + std::lock_guard g(mu_); + const int ln = l + 1; + for (int e = 0; e < NE; ++e) { + if (where_[(size_t) (ln * NE + e)] >= 0) continue; + const int s = victim(lay_[(size_t) ln].cls); + if (s < 0) break; + submit(s, ln, e, false); + slots_[(size_t) s].pin = stamp_; + ++p_streamed_; + any = true; + } + if (any) cv_.notify_all(); + } + (prompt ? p_host_ms_ : host_ms_) += now_ms() - t0; + return; + } + // read-ahead: layer l + 1 (each token's top pre_k_ of the predicted router logits), then layer l + 2 (top pre_k2_) + if (prefetch_ && pred && l + 1 < NL) ahead(l + 1, T, pred, pre_k_); + if (pre_k2_ > 0 && pred2 && l + 2 < NL) ahead(l + 2, T, pred2, pre_k2_); + host_ms_ += now_ms() - t0; +} + +// A region of `bytes` at the end of one arena buffer, for a prompt's scratch buffers: its slots are emptied and held +// (state 3) while a prompt runs (hold_region), and are ordinary slots again for decoding. +uint8_t* ExpertStream::region(uint64_t bytes) { + for (int c = (int) chunks_.size() - 1; c >= 0; --c) { + if (chunks_[(size_t) c].bytes < bytes + (64ull << 20)) continue; + uint8_t* end = chunks_[(size_t) c].as() + chunks_[(size_t) c].bytes; + uint8_t* lo = end - (bytes + PAGE - 1) / PAGE * PAGE; + region_slots_.clear(); + uint8_t* first = end; + for (int s = 0; s < (int) slots_.size(); ++s) { + const uint8_t* h = host_[(size_t) s]; + const uint64_t sz = cls_[(size_t) cls_of(s)].pages * PAGE; + if (h + sz > lo && h < end) { region_slots_.push_back(s); first = std::min(first, (uint8_t*) h); } + } + { // a class left with fewer slots than a prompt chunk's layer can use (NE) would stop with "no free slot": a small + // arena (a Mac with little free memory at start) gets the prompt's buffers as their own allocation instead + std::vector stolen(cls_.size(), 0); + for (int s : region_slots_) ++stolen[(size_t) cls_of(s)]; + bool tight = false; + for (size_t i = 0; i < cls_.size(); ++i) tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < NE); + if (tight) return standalone_region(bytes); + } + region_ = lo; + region_bytes_ = (uint64_t) (end - lo); + std::fprintf(stderr, "strata metal stream: %.2f GB of the arena (%zu slots) hold a prompt's buffers while it runs\n", + bytes / 1e9, region_slots_.size()); + return lo; + } + // no chunk is big enough for the prompt's buffers (an arena that shrank at the start, a 24 GB Mac): the fallback + // inside the loop above is only reached when a chunk fits, so the buffers get their own allocation here too + return standalone_region(bytes); +} + +// The prompt's buffers as their own anonymous allocation (fresh zero pages), when no region of the arena's slots +// takes them: no slots are held while a prompt runs, and the arena keeps decoding with all of its slots. +uint8_t* ExpertStream::standalone_region(uint64_t bytes) { + const uint64_t sz = (bytes + PAGE - 1) / PAGE * PAGE; + void* m = mmap(nullptr, (size_t) sz, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0); + if (m == MAP_FAILED) return nullptr; + region_slots_.clear(); + region_ = (uint8_t*) m; + region_bytes_ = 0; // fresh zero pages; nothing in the arena to empty + std::fprintf(stderr, "strata metal stream: the arena is small: the prompt's %.2f GB of buffers are their own " + "allocation, not a region of its slots\n", bytes / 1e9); + return region_; +} + +void ExpertStream::hold_region(bool hold) { + std::unique_lock g(mu_); + if (hold) { + // nothing may be reading into those slots: the queued read-aheads are dropped, the running ones waited for + std::vector drop; + for (const Job& j : pre_q_) + if (slots_[(size_t) j.slot].state.load() == 1 && std::find(drop.begin(), drop.end(), j.slot) == drop.end()) + drop.push_back(j.slot); + for (int s : drop) cancel_slot(s); + done_cv_.wait(g, [&] { return outstanding_ == 0; }); + for (int s : region_slots_) { + Slot& x = slots_[(size_t) s]; + if (x.layer >= 0) { + const int key = x.layer * NE + x.e; + where_[(size_t) key] = -1; + smap_.as()[key] = 0xFFFFFFFFu; + } + x.layer = x.e = -1; + x.pre = false; + x.state = 3; + } + if (region_) std::memset(region_, 0, (size_t) region_bytes_); + } else { + for (int s : region_slots_) { + slots_[(size_t) s].state = 0; + if (!retired_[(size_t) s]) lru_back(s); // a retired slot's memory serves the prompt only + } + } +} + +void ExpertStream::reset_stats() { + lookups_ = hits_ = pre_hits_ = late_ = misses_ = pre_issued_ = pre_cancel_ = windows_ = 0; + p_lookups_ = p_hits_ = p_pre_ = p_misses_ = p_streamed_ = p_chunks_ = 0; + p_wait_ms_ = p_host_ms_ = 0; + bytes_read_ = 0; + reads_ = 0; + wait_ms_ = host_ms_ = 0; +} + +void ExpertStream::report(FILE* f) const { + if (!ready()) return; + if (p_chunks_) + std::fprintf(f, "expert stream prompts %llu chunks: %llu distinct lookups (%.1f per layer), in the arena %.4f, streamed " + "ahead %.4f, missed %.4f; %llu streamed in; host %.1f ms per chunk, of it waiting %.1f ms\n", + (unsigned long long) p_chunks_, (unsigned long long) p_lookups_, p_lookups_ / 48.0 / p_chunks_, + p_hits_ / std::max(1.0, (double) p_lookups_), p_pre_ / std::max(1.0, (double) p_lookups_), + p_misses_ / std::max(1.0, (double) p_lookups_), (unsigned long long) p_streamed_, p_host_ms_ / p_chunks_, + p_wait_ms_ / p_chunks_); + const double w = std::max(1, windows_); + std::fprintf(f, + "expert stream %llu windows: %llu distinct lookups (%.1f per window), hits %.4f, read-ahead hits " + "%.4f (%llu of them late), misses %.4f (%.1f per window); read-aheads %llu (%.1f per window, %llu dropped); %.1f MB " + "read per window (%llu reads); host %.2f ms per window, of it waiting %.2f ms\n", + (unsigned long long) windows_, (unsigned long long) lookups_, lookups_ / w, + hits_ / std::max(1.0, (double) lookups_), pre_hits_ / std::max(1.0, (double) lookups_), + (unsigned long long) late_, misses_ / std::max(1.0, (double) lookups_), misses_ / w, + (unsigned long long) pre_issued_, pre_issued_ / w, (unsigned long long) pre_cancel_, + (bytes_read_.load() - 0) / 1e6 / w, (unsigned long long) reads_.load(), host_ms_ / w, wait_ms_ / w); +} + +} // namespace strata::metal_engine diff --git a/edge0-strata/metal/src/engine/stream.h b/edge0-strata/metal/src/engine/stream.h new file mode 100644 index 00000000..1a72af5f --- /dev/null +++ b/edge0-strata/metal/src/engine/stream.h @@ -0,0 +1,132 @@ +// metal/src/engine/stream.h - the SSD-streaming expert arena for Macs whose memory cannot hold the model (MAC-PLAN +// Phase 3, the small-Mac path; measured on the M4 Mac mini 32 GB, phaseA-tests/notes-mac-ssd.md). +// +// * The dense weights (everything but the routed experts) are read once into one resident shared buffer; the +// GGUF is never wrapped as a whole (a 37.6 GB NoCopy buffer cannot be made resident in 32 GB). +// * The routed experts live in a resident arena of slots (one shared MTLBuffer), filled straight from the GGUF +// with F_NOCACHE preads: one SSD read per expert slice, no staging copy, no page cache. smap (48 x 512 uint32: +// the slot's 16 KiB page in the arena, or ~0) is what the expert kernels read (moe.metal, function constant 20). +// * The decode window is one command buffer in 48 segments; after each layer's router the GPU signals a shared +// event and waits for the host (mtl_rt run_graph_seg). The host's turn (at_layer): the window's routed experts +// of that layer are looked up, the missing ones read into LRU victims (several reads in flight), smap updated; +// then the experts the NEXT layer's router predicts from this layer's FFN input are read ahead on the I/O +// threads while the GPU computes. +// * The arena starts with the expert profile's top-ranked pairs (data/expert-profile.bin). +// The numbers are exact: a slot holds the expert's GGUF bytes unchanged, so the answers are those of the resident +// engine (the kernels read the same bytes at another address). +#pragma once + +#include "gguf_lite.h" +#include "mtl_rt.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace strata::metal_engine { + +class ExpertStream { +public: + ~ExpertStream(); + // On when STRATA_METAL_STREAM=1, or (unset / auto) when the GGUF's experts do not fit in the GPU working set + // with room to spare. =0: never (the whole GGUF wrapped, as on the M5 Max). + static bool wanted(const gguf::File& f, const mtl::Context& cx); + bool open(const gguf::File& f, std::string& err); + // pread `len` bytes at file offset `off` (F_NOCACHE) + bool read(void* dst, uint64_t len, uint64_t off) const; + // The dense tensors into one buffer; every tensor's offset in `f` is rewritten to its offset in `out`. + // Tensors named in `skip` are left out (their offset is set to 0: they must not be read any more). + bool load_dense(gguf::File& f, mtl::Context& cx, const std::vector& skip, mtl::Buffer& out, + std::string& err); + // The arena: `budget` bytes of slots, filled from the profile's ranking (or layer-interleaved if none). + // min_slots: each size class at least this many slots (576 with prompt chunks, else 64) + // recheck (optional): just before the arena is committed, the bytes it may have now (arena_shrink); when that is + // less than the arena laid out, nothing is allocated, shrink_to() says how much, and the caller tries again + bool make_arena(mtl::Context& cx, uint64_t budget, const std::string& profile, int min_slots, std::string& err, + const std::function& recheck = {}, + const std::function& grow = {}); + uint64_t shrink_to() const { return shrink_to_; } + uint64_t arena_used_bytes() const { return arena_live_; } // the arena's bytes in use (its retired slots left out) + bool ready() const { return !chunks_.empty(); } + // the arena's buffers (at most 4, each within the device's maxBufferLength): chunk(i) for i >= count is chunk 0 + const mtl::Buffer& chunk(int i) const { return chunks_[(size_t) i < chunks_.size() ? (size_t) i : 0]; } + const mtl::Buffer& smap() const { return smap_; } + uint64_t up_in_slot(int l) const { return lay_[(size_t) l].gs; } + uint64_t down_in_slot(int l) const { return lay_[(size_t) l].gs + lay_[(size_t) l].us; } + // The host's turn at layer l: ids = the window's T x topk routed ids; pred = T x 512 router logits of layer l + 1 + // computed from this layer's FFN input (nullptr: no read-ahead). + // pred2: layer l + 2's, from the same input (STRATA_METAL_PREFETCH2 = its top-k, 0 = off) + void at_layer(int l, int T, int topk, const int32_t* ids, const float* pred, const float* pred2 = nullptr); + int ahead2() const { return pre_k2_; } + // a prompt's scratch buffers inside the arena: region() places them (once), hold_region() empties and holds its + // slots while a prompt runs and gives them back after + uint8_t* region(uint64_t bytes); + void hold_region(bool hold); + void report(FILE* f) const; + void reset_stats(); + +private: + uint64_t shrink_to_ = 0; + struct Lay { uint64_t g, gs, u, us, d, ds; int cls; }; + struct Slot { int16_t layer = -1; int16_t e = -1; int prev = -1, next = -1; bool pre = false; int left = 0; uint64_t pin = 0; std::atomic state{0}; }; // 0 idle, 1 queued, 2 reading + struct Cls { uint64_t pages = 0; int n = 0; int base = 0; int head = -1, tail = -1; }; // slot range [base, base+n) + struct Job { int slot; int layer, e, part; }; // part: slice * split_ + piece + void lru_unlink(int s); + void lru_front(int s); + void lru_back(int s); + int cls_of(int s) const; + int victim(int cls); + void read_part(int slot, int l, int e, int part); + void cancel_slot(int s); + void promote_slot(int s); + void io_main(); + void ahead(int ln, int T, const float* pred, int K); + void submit(int slot, int l, int e, bool urgent); + void wait_slot(int s); + uint8_t* standalone_region(uint64_t bytes); // the prompt's buffers as their own allocation, when no region of the arena's slots takes them + + int fd_ = -1; + // .experts (metal/tools/stream_experts.py): each expert's gate, up and down rows together, read with one + // request (cfd_ >= 0); else the GGUF's three slices. STRATA_METAL_CONTIG=0: the GGUF always. + int cfd_ = -1; + std::vector coff_; + std::vector lay_; + std::vector cls_; + std::vector slots_; + std::vector retired_; // slots left unused when the stepped fill stopped at macOS's floor (never in an LRU) + std::vector where_; // (l * 512 + e) -> slot, -1 + std::vector chunks_; + std::vector region_slots_; + uint8_t* region_ = nullptr; + uint64_t region_bytes_ = 0; + std::vector code_; // slot -> its smap value (chunk << 28 | page) + std::vector host_; // slot -> its CPU address + mtl::Buffer smap_; + uint64_t arena_bytes_ = 0; + uint64_t arena_live_ = 0; + // the I/O threads + std::vector io_; + std::mutex mu_; + std::condition_variable cv_, done_cv_; + std::deque urgent_q_, pre_q_; + bool stop_ = false; + int outstanding_ = 0; + int nio_ = 6, pre_k_ = 10, pre_k2_ = 0, split_ = 2; + bool prefetch_ = true, fake_ = false, pre_mru_ = false; + // stats + uint64_t stamp_ = 1; // the handshake count: slots pinned with the current one are in use on the GPU + uint64_t p_lookups_ = 0, p_hits_ = 0, p_pre_ = 0, p_misses_ = 0, p_streamed_ = 0, p_chunks_ = 0; + double p_wait_ms_ = 0, p_host_ms_ = 0; + uint64_t lookups_ = 0, hits_ = 0, pre_hits_ = 0, late_ = 0, misses_ = 0, pre_issued_ = 0, pre_cancel_ = 0, windows_ = 0; + std::atomic bytes_read_{0}, reads_{0}; + double wait_ms_ = 0, host_ms_ = 0, fill_s_ = 0; + uint64_t fill_bytes_ = 0; +}; + +} // namespace strata::metal_engine diff --git a/edge0-strata/metal/src/engine/stubs.cpp b/edge0-strata/metal/src/engine/stubs.cpp new file mode 100644 index 00000000..7f8f0a9a --- /dev/null +++ b/edge0-strata/metal/src/engine/stubs.cpp @@ -0,0 +1,371 @@ +// metal/src/engine/stubs.cpp - GENERATED by metal/tools/gen_stubs.py: do not edit. +// The kernel-library entry points with no Metal implementation yet: each stops the engine with its name +// (metal/src/engine/missing.cpp). Regenerate after porting one. + +#include "strata/kernels/bf16_bits.hpp" +#include "strata/kernels/bf16_gemv.hpp" +#include "strata/kernels/cvec.hpp" +#include "strata/kernels/dequant_bf16.hpp" +#include "strata/kernels/dequant_s2.hpp" +#include "strata/kernels/elementwise.hpp" +#include "strata/kernels/f16_bits.hpp" +#include "strata/kernels/fused_gdn.hpp" +#include "strata/kernels/fused_gr.hpp" +#include "strata/kernels/gdn.hpp" +#include "strata/kernels/gfx_arch.hpp" +#include "strata/kernels/gr.hpp" +#include "strata/kernels/iq_kernels.hpp" +#include "strata/kernels/kv_q4.hpp" +#include "strata/kernels/kv_q8.hpp" +#include "strata/kernels/kv_stream.hpp" +#include "strata/kernels/mrope.hpp" +#include "strata/kernels/native_flash_attn.hpp" +#include "strata/kernels/native_gdn.hpp" +#include "strata/kernels/native_gdn_preprocess.hpp" +#include "strata/kernels/native_gr_norm.hpp" +#include "strata/kernels/native_gr_postops.hpp" +#include "strata/kernels/native_mmvq.hpp" +#include "strata/kernels/native_moe.hpp" +#include "strata/kernels/native_ple_postops.hpp" +#include "strata/kernels/native_qsa.hpp" +#include "strata/kernels/native_qsa_indexer.hpp" +#include "strata/kernels/native_qsa_score.hpp" +#include "strata/kernels/native_rope.hpp" +#include "strata/kernels/native_router.hpp" +#include "strata/kernels/ngram.hpp" +#include "strata/kernels/pdl.hpp" +#include "strata/kernels/ple.hpp" +#include "strata/kernels/qsa.hpp" +#include "strata/kernels/qsa_decode_attn.hpp" +#include "strata/kernels/qsa_prompt_attn.hpp" +#include "strata/kernels/qsa_select.hpp" +#include "strata/kernels/quantize_act.hpp" +#include "strata/kernels/rope.hpp" +#include "strata/kernels/rope_scaling.hpp" +#include "strata/kernels/route_prior.hpp" +#include "strata/kernels/router_top10.hpp" +#include "strata/kernels/s2_expert_grouped.hpp" +#include "strata/kernels/s2_gemv.hpp" +#include "strata/kernels/s2_gemv_q8.hpp" +#include "strata/kernels/s_gemv.hpp" +#include "strata/kernels/sampler.hpp" +#include "strata/kernels/shared_expert.hpp" +#include "strata/kernels/verify_kernels.hpp" +#include "strata/kernels/cpu/expert.hpp" +#include "strata/kernels/cpu/expert_layout.hpp" +#include "strata/kernels/cpu/iq_avx1.hpp" +#include "strata/kernels/cpu/iq_avx2.hpp" +#include "strata/kernels/cpu/iq_avx512.hpp" +#include "strata/kernels/cpu/kq_avx1.hpp" +#include "strata/kernels/cpu/kq_avx2.hpp" +#include "strata/kernels/cpu/native_expert.hpp" +#include "strata/kernels/cpu/pool.hpp" +#include "strata/prefill/gemm.hpp" +#include "strata/prefill/kernels.hpp" +#include "strata/prefill/moe_fused.hpp" +#include "strata/prefill/moe_fused_iq.hpp" +#include "strata/prefill/moe_mmq.hpp" +#include "strata/prefill/prefill.hpp" +#include "strata/prefill/share_rules.hpp" +#include "strata/prefill/xmx_moe.hpp" +#include "strata/core/arch_defaults.hpp" +#include "strata/core/batch_rows.hpp" +#include "strata/core/conversation_buffer.hpp" +#include "strata/core/conversation_cache.hpp" +#include "strata/core/conversation_file.hpp" +#include "strata/core/conversation_memory.hpp" +#include "strata/core/conversation_snapshot.hpp" +#include "strata/core/coupled_draft.hpp" +#include "strata/core/device.hpp" +#include "strata/core/dma_batch.hpp" +#include "strata/core/emulate.hpp" +#include "strata/core/exchange_storage.hpp" +#include "strata/core/expert_cache.hpp" +#include "strata/core/expert_source.hpp" +#include "strata/core/foresight_swap.hpp" +#include "strata/core/graph.hpp" +#include "strata/core/hit_hook.hpp" +#include "strata/core/layer.hpp" +#include "strata/core/layout.hpp" +#include "strata/core/logit_bias.hpp" +#include "strata/core/mtp.hpp" +#include "strata/core/native_dense.hpp" +#include "strata/core/native_head.hpp" +#include "strata/core/on_device.hpp" +#include "strata/core/peer_experts.hpp" +#include "strata/core/pinned.hpp" +#include "strata/core/pinned_ring.hpp" +#include "strata/core/progress.hpp" +#include "strata/core/remote_expert_opt.hpp" +#include "strata/core/remote_experts.hpp" +#include "strata/core/session.hpp" +#include "strata/core/session_save_reclaim.hpp" +#include "strata/core/spec_prob.hpp" +#include "strata/core/verify.hpp" +#include "strata/core/vmm.hpp" +#include "strata/core/weights.hpp" +#include "strata/platform/aux_cpus.hpp" +#include "strata/platform/direct_file.hpp" +#include "strata/platform/host_limits.hpp" +#include "strata/platform/memory.hpp" +#include "strata/ngram/ple_reader.hpp" +#include "strata/spec/controller.hpp" +#include "strata/spec/draft_policy.hpp" +#include "strata/spec/draft_source.hpp" +#include "strata/spec/suffix_drafter.hpp" +#include + +namespace strata::metal_engine { void missing(const char* name); } + +auto strata::core::RemoteExpertOpt::accumulate(const RemoteExperts & p0) -> void { strata::metal_engine::missing("strata::core::RemoteExpertOpt::accumulate"); } +auto strata::core::RemoteExpertOpt::adapt(const std::vector & p0, const std::vector & p1, const std::vector> & p2, int p3, ExpertSource & p4) -> bool { strata::metal_engine::missing("strata::core::RemoteExpertOpt::adapt"); return {}; } +auto strata::core::RemoteExpertOpt::attach(RemoteExperts & p0) -> void { strata::metal_engine::missing("strata::core::RemoteExpertOpt::attach"); } +auto strata::core::RemoteExpertOpt::begin(const float * p0, int p1, int p2) -> void { strata::metal_engine::missing("strata::core::RemoteExpertOpt::begin"); } +auto strata::core::RemoteExpertOpt::combine(float * p0, float * p1, int p2, int p3, const uint32_t * p4, uint32_t p5, void * p6) const -> void { strata::metal_engine::missing("strata::core::RemoteExpertOpt::combine"); } +auto strata::core::RemoteExpertOpt::copy_rows(float * p0, const float * p1, int p2, int p3, const int32_t * p4, const int32_t * p5, void * p6) const -> void { strata::metal_engine::missing("strata::core::RemoteExpertOpt::copy_rows"); } +auto strata::core::RemoteExpertOpt::init(std::string & p0) -> bool { strata::metal_engine::missing("strata::core::RemoteExpertOpt::init"); return {}; } +auto strata::core::RemoteExpertOpt::metadata_bytes() -> size_t { strata::metal_engine::missing("strata::core::RemoteExpertOpt::metadata_bytes"); return {}; } +auto strata::core::RemoteExpertOpt::owns(int64_t p0, int32_t p1) const -> bool { strata::metal_engine::missing("strata::core::RemoteExpertOpt::owns"); return {}; } +auto strata::core::RemoteExpertOpt::prepare(const RemoteExperts & p0, void * p1) const -> void { strata::metal_engine::missing("strata::core::RemoteExpertOpt::prepare"); } +auto strata::core::RemoteExpertOpt::reduce(RemoteExperts & p0, const void * p1, std::string & p2) -> bool { strata::metal_engine::missing("strata::core::RemoteExpertOpt::reduce"); return {}; } +strata::core::RemoteExpertOpt::~RemoteExpertOpt() { strata::metal_engine::missing("strata::core::RemoteExpertOpt::~RemoteExpertOpt"); } +namespace strata::kernels { void add_inplace(float * p0, const float * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::kernels::add_inplace"); } } +namespace strata::kernels { void add_streams_broadcast(const float * p0, const float * p1, float * p2, int64_t p3, int p4, int p5, void * p6) { strata::metal_engine::missing("strata::kernels::add_streams_broadcast"); } } +namespace strata::kernels { void argmax_rows(const float * p0, int p1, int p2, void * p3, int32_t * p4, void * p5) { strata::metal_engine::missing("strata::kernels::argmax_rows"); } } +namespace strata::kernels { bool argmax_rows_wanted() { strata::metal_engine::missing("strata::kernels::argmax_rows_wanted"); return {}; } } +namespace strata::kernels { void bf16_gemv(const uint16_t * p0, const uint16_t * p1, float * p2, int64_t p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::kernels::bf16_gemv"); } } +namespace strata::kernels { void bf16_gemv_fp32_mmvf(const float * p0, const uint16_t * p1, float * p2, int64_t p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::kernels::bf16_gemv_fp32_mmvf"); } } +namespace strata::kernels { void bf16_gemv_fp32_mmvf_cols(const float * p0, const uint16_t * p1, float * p2, int64_t p3, int64_t p4, int p5, void * p6) { strata::metal_engine::missing("strata::kernels::bf16_gemv_fp32_mmvf_cols"); } } +namespace strata::kernels { void bf16_gemv_fp32_mmvf_multi(const float * p0, int64_t p1, const uint16_t * p2, float * p3, int64_t p4, int64_t p5, int64_t p6, int p7, void * p8) { strata::metal_engine::missing("strata::kernels::bf16_gemv_fp32_mmvf_multi"); } } +namespace strata::kernels { bool bf16_gemv_fp32_mmvf_multi_aux(const float * p0, int64_t p1, const uint16_t * p2, float * p3, int64_t p4, int64_t p5, int64_t p6, int p7, const uint16_t * p8, float * p9, int64_t p10, void * p11) { strata::metal_engine::missing("strata::kernels::bf16_gemv_fp32_mmvf_multi_aux"); return {}; } } +namespace strata::kernels { void bf16_gemv_split(const uint16_t * p0, const uint16_t * p1, float * p2, int64_t p3, int64_t p4, int p5, void * p6) { strata::metal_engine::missing("strata::kernels::bf16_gemv_split"); } } +namespace strata::kernels { void broadcast_streams(const float * p0, float * p1, int64_t p2, int p3, int p4, void * p5) { strata::metal_engine::missing("strata::kernels::broadcast_streams"); } } +namespace strata::kernels { void copy_from_mapped(float * p0, const float * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::kernels::copy_from_mapped"); } } +namespace strata::kernels { void copy_from_mapped_multi(const MappedCopy * p0, int p1, void * p2) { strata::metal_engine::missing("strata::kernels::copy_from_mapped_multi"); } } +namespace strata::kernels { void copy_i32_from_mapped(int32_t * p0, const int32_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::kernels::copy_i32_from_mapped"); } } +namespace strata::kernels { void copy_i32_from_mapped_unless(int32_t * p0, const int32_t * p1, long long p2, const uint32_t * p3, uint32_t p4, void * p5) { strata::metal_engine::missing("strata::kernels::copy_i32_from_mapped_unless"); } } +namespace strata::kernels { void copy_indexed(float * p0, const float * p1, int64_t p2, const int32_t * p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::kernels::copy_indexed"); } } +namespace strata::kernels { void copy_or_zero_from_mapped(float * p0, const float * p1, long long p2, const uint32_t * p3, uint32_t p4, void * p5) { strata::metal_engine::missing("strata::kernels::copy_or_zero_from_mapped"); } } +namespace strata::kernels { void copy_row_to_first(const int32_t * p0, float * p1, int64_t p2, float * p3, int64_t p4, float * p5, int64_t p6, void * p7) { strata::metal_engine::missing("strata::kernels::copy_row_to_first"); } } +namespace strata::kernels { void copy_rows_from_mapped(float * p0, const float * p1, int64_t p2, int64_t p3, const int32_t * p4, const int32_t * p5, void * p6) { strata::metal_engine::missing("strata::kernels::copy_rows_from_mapped"); } } +namespace strata::kernels { void copy_rows_strided(float * p0, const float * p1, int64_t p2, int64_t p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::kernels::copy_rows_strided"); } } +namespace strata::kernels { void coupled_draft_sample(float * p0, int p1, const int32_t * p2, const int32_t * p3, int p4, const SamplerParams * p5, int32_t * p6, int p7, int p8, const int32_t * p9, void * p10, int32_t * p11, float * p12, void * p13) { strata::metal_engine::missing("strata::kernels::coupled_draft_sample"); } } +namespace strata::kernels { size_t coupled_draft_scratch_bytes(int p0) { strata::metal_engine::missing("strata::kernels::coupled_draft_scratch_bytes"); return {}; } } +namespace strata::kernels { void coupled_draft_stage(const SamplerParams * p0, const int32_t * p1, SamplerParams * p2, int32_t * p3, int p4, void * p5) { strata::metal_engine::missing("strata::kernels::coupled_draft_stage"); } } +namespace strata::kernels::cpu { void act_quant_q8_1(const float * p0, int p1, ActQ & p2) { strata::metal_engine::missing("strata::kernels::cpu::act_quant_q8_1"); } } +namespace strata::kernels::cpu { void act_quant_q8_1_avx1(const float * p0, int p1, ActQ & p2) { strata::metal_engine::missing("strata::kernels::cpu::act_quant_q8_1_avx1"); } } +namespace strata::kernels::cpu { void act_quant_q8_1_avx2(const float * p0, int p1, ActQ & p2) { strata::metal_engine::missing("strata::kernels::cpu::act_quant_q8_1_avx2"); } } +namespace strata::kernels::cpu { void bf16_rows_dot_multi(const uint16_t * p0, int p1, int p2, const float * p3, int p4, float * p5) { strata::metal_engine::missing("strata::kernels::cpu::bf16_rows_dot_multi"); } } +namespace strata::kernels::cpu { void bf16_rows_dot_multi_avx1(const uint16_t * p0, int p1, int p2, const float * p3, int p4, float * p5) { strata::metal_engine::missing("strata::kernels::cpu::bf16_rows_dot_multi_avx1"); } } +namespace strata::kernels::cpu { void iq128_down_rows(int p0, const uint8_t * p1, size_t p2, int p3, const void *const * p4, int p5, float *const * p6, int p7, int p8) { strata::metal_engine::missing("strata::kernels::cpu::iq128_down_rows"); } } +namespace strata::kernels::cpu { bool iq128_down_supported(int p0) noexcept { strata::metal_engine::missing("strata::kernels::cpu::iq128_down_supported"); return {}; } } +namespace strata::kernels::cpu { void iq128_gu_rows(int p0, const uint8_t * p1, size_t p2, size_t p3, int p4, const void *const * p5, int p6, float *const * p7, int p8, int p9) { strata::metal_engine::missing("strata::kernels::cpu::iq128_gu_rows"); } } +namespace strata::kernels::cpu { bool iq128_supported(int p0) noexcept { strata::metal_engine::missing("strata::kernels::cpu::iq128_supported"); return {}; } } +namespace strata::kernels::cpu { void iq256_gu_rows(int p0, const uint8_t * p1, size_t p2, size_t p3, int p4, const void *const * p5, int p6, float *const * p7, int p8, int p9) { strata::metal_engine::missing("strata::kernels::cpu::iq256_gu_rows"); } } +namespace strata::kernels::cpu { void iq4nl256_down_rows(const uint8_t * p0, size_t p1, int p2, const void *const * p3, int p4, float *const * p5, int p6, int p7) { strata::metal_engine::missing("strata::kernels::cpu::iq4nl256_down_rows"); } } +namespace strata::kernels::cpu { void iq512_gu_rows(int p0, const uint8_t * p1, size_t p2, size_t p3, int p4, const void *const * p5, int p6, float *const * p7, int p8, int p9) { strata::metal_engine::missing("strata::kernels::cpu::iq512_gu_rows"); } } +namespace strata::kernels::cpu { void kq256_gu_rows(int p0, const uint8_t * p1, size_t p2, size_t p3, int p4, const void *const * p5, int p6, float *const * p7, int p8, int p9) { strata::metal_engine::missing("strata::kernels::cpu::kq256_gu_rows"); } } +namespace strata::kernels::cpu { void kq256_rows(int p0, const uint8_t * p1, size_t p2, int p3, const void *const * p4, int p5, float *const * p6, int p7, int p8) { strata::metal_engine::missing("strata::kernels::cpu::kq256_rows"); } } +namespace strata::kernels::cpu { void q2_0_gguf_rows_multi(const uint8_t * p0, size_t p1, int p2, const ActQ *const * p3, int p4, float *const * p5, int p6, int p7) { strata::metal_engine::missing("strata::kernels::cpu::q2_0_gguf_rows_multi"); } } +namespace strata::kernels::cpu { void q2_0_gguf_rows_multi_avx1(const uint8_t * p0, size_t p1, int p2, const ActQ *const * p3, int p4, float *const * p5, int p6, int p7) { strata::metal_engine::missing("strata::kernels::cpu::q2_0_gguf_rows_multi_avx1"); } } +namespace strata::kernels::cpu { void q2_0_gguf_rows_multi_avx1_legacy(const uint8_t * p0, size_t p1, int p2, const ActQ *const * p3, int p4, float *const * p5, int p6, int p7) { strata::metal_engine::missing("strata::kernels::cpu::q2_0_gguf_rows_multi_avx1_legacy"); } } +namespace strata::kernels::cpu { void q2_0_gguf_rows_multi_avx2(const uint8_t * p0, size_t p1, int p2, const ActQ *const * p3, int p4, float *const * p5, int p6, int p7) { strata::metal_engine::missing("strata::kernels::cpu::q2_0_gguf_rows_multi_avx2"); } } +namespace strata::kernels::cpu { void q2_0_gguf_rows_multi_avx2_legacy(const uint8_t * p0, size_t p1, int p2, const ActQ *const * p3, int p4, float *const * p5, int p6, int p7) { strata::metal_engine::missing("strata::kernels::cpu::q2_0_gguf_rows_multi_avx2_legacy"); } } +namespace strata::kernels::cpu { void q8k_quant_avx1(const float * p0, void * p1, int64_t p2) { strata::metal_engine::missing("strata::kernels::cpu::q8k_quant_avx1"); } } +namespace strata::kernels::cpu { void q8k_quant_avx2(const float * p0, void * p1, int64_t p2) { strata::metal_engine::missing("strata::kernels::cpu::q8k_quant_avx2"); } } +namespace strata::kernels::cpu { void s2_expert_down_rows(const uint8_t * p0, const ActQ & p1, float * p2, int p3, int p4) { strata::metal_engine::missing("strata::kernels::cpu::s2_expert_down_rows"); } } +namespace strata::kernels::cpu { void s2_expert_down_rows_multi(const uint8_t * p0, const ActQ *const * p1, int p2, float *const * p3, int p4, int p5) { strata::metal_engine::missing("strata::kernels::cpu::s2_expert_down_rows_multi"); } } +namespace strata::kernels::cpu { void s2_expert_gu_rows(const uint8_t * p0, const ActQ & p1, float * p2, int p3, int p4) { strata::metal_engine::missing("strata::kernels::cpu::s2_expert_gu_rows"); } } +namespace strata::kernels::cpu { void s2_expert_gu_rows_multi(const uint8_t * p0, const ActQ *const * p1, int p2, float *const * p3, int p4, int p5) { strata::metal_engine::missing("strata::kernels::cpu::s2_expert_gu_rows_multi"); } } +namespace strata::kernels::cpu { void s2_expert_vnni_q(const uint8_t * p0, const ActQ & p1, float * p2, ExpertScratch & p3) { strata::metal_engine::missing("strata::kernels::cpu::s2_expert_vnni_q"); } } +namespace strata::kernels { void cvec_apply(float * p0, int64_t p1, int64_t p2, int64_t p3, const float * p4, int64_t p5, const float * p6, int64_t p7, bool p8, void * p9) { strata::metal_engine::missing("strata::kernels::cvec_apply"); } } +namespace strata::kernels { bool cvec_replicate(std::string & p0) { strata::metal_engine::missing("strata::kernels::cvec_replicate"); return {}; } } +namespace strata::kernels { void cvec_set_enabled(bool p0) { strata::metal_engine::missing("strata::kernels::cvec_set_enabled"); } } +namespace strata::kernels { bool cvec_tables(const float ** p0, const float ** p1, const int ** p2) { strata::metal_engine::missing("strata::kernels::cvec_tables"); return {}; } } +namespace strata::kernels { bool cvec_upload(const std::vector & p0, const std::vector & p1, int p2, int p3, int p4, int64_t p5, int64_t p6, std::string & p7) { strata::metal_engine::missing("strata::kernels::cvec_upload"); return {}; } } +namespace strata::kernels { void doorbell_publish(const float * p0, const int32_t * p1, const float * p2, int64_t p3, int64_t p4, float * p5, int32_t * p6, float * p7, uint32_t * p8, void * p9) { strata::metal_engine::missing("strata::kernels::doorbell_publish"); } } +namespace strata::kernels { void doorbell_publish_res(const float * p0, const int32_t * p1, const int32_t * p2, int p3, int64_t p4, int64_t p5, float * p6, int32_t * p7, uint32_t * p8, void * p9) { strata::metal_engine::missing("strata::kernels::doorbell_publish_res"); } } +namespace strata::kernels { void doorbell_ring(uint32_t * p0, void * p1) { strata::metal_engine::missing("strata::kernels::doorbell_ring"); } } +namespace strata::kernels { void doorbell_wait(const uint32_t * p0, const uint32_t * p1, void * p2) { strata::metal_engine::missing("strata::kernels::doorbell_wait"); } } +namespace strata::kernels { void embedding_gather(const uint8_t * p0, const float * p1, const float * p2, int64_t p3, int p4, int p5, int p6, float * p7, void * p8) { strata::metal_engine::missing("strata::kernels::embedding_gather"); } } +namespace strata::kernels { void embedding_gather_dev(const uint8_t * p0, const float * p1, const float * p2, const int32_t * p3, int p4, int64_t p5, int p6, int p7, int p8, uint64_t p9, uint64_t p10, float * p11, void * p12) { strata::metal_engine::missing("strata::kernels::embedding_gather_dev"); } } +namespace strata::kernels { void f32_to_bf16_bulk(const float * p0, uint16_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::kernels::f32_to_bf16_bulk"); } } +namespace strata::kernels { void fetch_blobs(const unsigned long long * p0, const int32_t * p1, uint8_t * p2, int64_t p3, int p4, void * p5) { strata::metal_engine::missing("strata::kernels::fetch_blobs"); } } +namespace strata::kernels { void force_token(int32_t * p0, const int32_t * p1, int p2, void * p3) { strata::metal_engine::missing("strata::kernels::force_token"); } } +namespace strata::kernels { void fused_gdn_ab(const float * p0, const uint16_t * p1, const uint16_t * p2, const float * p3, const float * p4, float * p5, float * p6, int p7, int p8, void * p9) { strata::metal_engine::missing("strata::kernels::fused_gdn_ab"); } } +namespace strata::kernels { void fused_gdn_conv_l2(float * p0, const float * p1, const float * p2, float * p3, int p4, int p5, float p6, void * p7) { strata::metal_engine::missing("strata::kernels::fused_gdn_conv_l2"); } } +namespace strata::kernels { void fused_gdn_step_norm(float * p0, const float * p1, const float * p2, const float * p3, const float * p4, const float * p5, const float * p6, const float * p7, float p8, float * p9, int p10, int p11, void * p12) { strata::metal_engine::missing("strata::kernels::fused_gdn_step_norm"); } } +namespace strata::kernels { void fused_gr_dequant_q8_0(const uint8_t * p0, uint16_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::kernels::fused_gr_dequant_q8_0"); } } +namespace strata::kernels { void fused_gr_quantize_q8_0(const uint16_t * p0, uint8_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::kernels::fused_gr_quantize_q8_0"); } } +namespace strata::kernels { void fused_gr_read(const FusedGrArgs & p0, void * p1) { strata::metal_engine::missing("strata::kernels::fused_gr_read"); } } +namespace strata::kernels { bool fused_gr_read_multi(const FusedGrArgs * p0, int p1, float * p2, void * p3, unsigned long long * p4, int p5) { strata::metal_engine::missing("strata::kernels::fused_gr_read_multi"); return {}; } } +namespace strata::kernels { bool fused_gr_supported(int64_t p0, int64_t p1, int64_t p2) { strata::metal_engine::missing("strata::kernels::fused_gr_supported"); return {}; } } +namespace strata::kernels { void fwht256_cuda(const float * p0, float * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::kernels::fwht256_cuda"); } } +namespace strata::kernels { void gdn_ab_multi(const float * p0, const uint16_t * p1, const uint16_t * p2, const float * p3, const float * p4, float * p5, float * p6, int p7, int p8, int p9, void * p10) { strata::metal_engine::missing("strata::kernels::gdn_ab_multi"); } } +namespace strata::kernels { void gdn_beta_gate(float * p0, int64_t p1, void * p2) { strata::metal_engine::missing("strata::kernels::gdn_beta_gate"); } } +namespace strata::kernels { void gdn_conv_commit(float * p0, const float * p1, int p2, const int32_t * p3, void * p4) { strata::metal_engine::missing("strata::kernels::gdn_conv_commit"); } } +namespace strata::kernels { void gdn_conv_l2_multi(const float * p0, const float * p1, const float * p2, float * p3, int p4, int p5, float p6, int p7, void * p8, int p9, bool p10) { strata::metal_engine::missing("strata::kernels::gdn_conv_l2_multi"); } } +namespace strata::kernels { void gdn_conv_step(float * p0, const float * p1, const float * p2, float * p3, int64_t p4, int64_t p5, void * p6) { strata::metal_engine::missing("strata::kernels::gdn_conv_step"); } } +namespace strata::kernels { void gdn_gate(const float * p0, const float * p1, const float * p2, float * p3, int64_t p4, int64_t p5, void * p6) { strata::metal_engine::missing("strata::kernels::gdn_gate"); } } +namespace strata::kernels { void gdn_l2_norm(float * p0, int64_t p1, int64_t p2, float p3, void * p4) { strata::metal_engine::missing("strata::kernels::gdn_l2_norm"); } } +namespace strata::kernels { void gdn_out_norm(const float * p0, const float * p1, const float * p2, float * p3, int64_t p4, int64_t p5, float p6, void * p7) { strata::metal_engine::missing("strata::kernels::gdn_out_norm"); } } +namespace strata::kernels { void gdn_step(float * p0, const float * p1, const float * p2, const float * p3, const float * p4, const float * p5, float * p6, const GdnShapes & p7, void * p8) { strata::metal_engine::missing("strata::kernels::gdn_step"); } } +namespace strata::kernels { void gdn_step_norm_multi(float * p0, const float * p1, int p2, const float * p3, const float * p4, const float * p5, const float * p6, float p7, float * p8, int p9, int p10, int p11, const int32_t * p12, void * p13, int p14, void * p15) { strata::metal_engine::missing("strata::kernels::gdn_step_norm_multi"); } } +namespace strata::kernels { void gpu_noop(void * p0) { strata::metal_engine::missing("strata::kernels::gpu_noop"); } } +namespace strata::kernels { void gpu_stamp(unsigned long long * p0, int p1, void * p2) { strata::metal_engine::missing("strata::kernels::gpu_stamp"); } } +namespace strata::kernels { void gr_read(const float * p0, const float * p1, const uint16_t * p2, const uint16_t * p3, const uint16_t * p4, float p5, const GrShapes & p6, const GrWorkspace & p7, float * p8, float * p9, void * p10) { strata::metal_engine::missing("strata::kernels::gr_read"); } } +namespace strata::kernels { void gr_write(const float * p0, const float * p1, const float * p2, const GrShapes & p3, float * p4, void * p5) { strata::metal_engine::missing("strata::kernels::gr_write"); } } +namespace strata::kernels { void gr_write_multi(const float * p0, const float * p1, const float * p2, const GrShapes & p3, float * p4, int p5, void * p6) { strata::metal_engine::missing("strata::kernels::gr_write_multi"); } } +namespace strata::kernels { void indexer_key_append(const float * p0, const int32_t * p1, int32_t p2, const float * p3, float p4, const QsaIndexerBuffers & p5, const QsaShapes & p6, const float * p7, const float * p8, void * p9) { strata::metal_engine::missing("strata::kernels::indexer_key_append"); } } +namespace strata::kernels { void iq_dequant_f16(int p0, const void * p1, int64_t p2, uint16_t * p3, void * p4) { strata::metal_engine::missing("strata::kernels::iq_dequant_f16"); } } +namespace strata::kernels { void iq_dequant_f32(int p0, const void * p1, int64_t p2, float * p3, void * p4) { strata::metal_engine::missing("strata::kernels::iq_dequant_f32"); } } +namespace strata::kernels { void iq_dequant_gu_f16(int p0, const void * p1, const void * p2, int64_t p3, int64_t p4, uint16_t * p5, void * p6) { strata::metal_engine::missing("strata::kernels::iq_dequant_gu_f16"); } } +namespace strata::kernels { void iq_embed_rows(int p0, const void * p1, size_t p2, const int32_t * p3, int64_t p4, int64_t p5, float * p6, void * p7) { strata::metal_engine::missing("strata::kernels::iq_embed_rows"); } } +namespace strata::kernels { void kv_append_q4(uint8_t * p0, uint8_t * p1, const int32_t * p2, int64_t p3, int64_t p4, const float * p5, const float * p6, const QsaShapes & p7, void * p8, const KvHostPools * p9, const KvHostPools * p10) { strata::metal_engine::missing("strata::kernels::kv_append_q4"); } } +namespace strata::kernels { void kv_append_q4_step(uint8_t * p0, uint8_t * p1, const int32_t * p2, const int32_t * p3, const float * p4, const float * p5, const QsaShapes & p6, void * p7, const KvHostPools * p8) { strata::metal_engine::missing("strata::kernels::kv_append_q4_step"); } } +namespace strata::kernels { void kv_append_q4_steps(uint8_t * p0, uint8_t * p1, const int32_t * p2, const int32_t * p3, int p4, int p5, const float * p6, const float * p7, const QsaShapes & p8, void * p9, const KvHostPools * p10) { strata::metal_engine::missing("strata::kernels::kv_append_q4_steps"); } } +namespace strata::kernels { void kv_append_q8_step(int8_t * p0, int8_t * p1, uint16_t * p2, uint16_t * p3, const int32_t * p4, const int32_t * p5, const float * p6, const float * p7, const QsaShapes & p8, void * p9, const KvHostPools * p10) { strata::metal_engine::missing("strata::kernels::kv_append_q8_step"); } } +namespace strata::kernels { void kv_append_q8_steps(int8_t * p0, int8_t * p1, uint16_t * p2, uint16_t * p3, const int32_t * p4, const int32_t * p5, int p6, const float * p7, const float * p8, int p9, int p10, const QsaShapes & p11, void * p12, const KvHostPools * p13) { strata::metal_engine::missing("strata::kernels::kv_append_q8_steps"); } } +namespace strata::kernels { void kv_append_step(uint16_t * p0, uint16_t * p1, const int32_t * p2, const int32_t * p3, const float * p4, const float * p5, const QsaShapes & p6, void * p7, const KvHostPools * p8) { strata::metal_engine::missing("strata::kernels::kv_append_step"); } } +namespace strata::kernels { uint64_t kv_block_bytes(const QsaShapes & p0, int p1) { strata::metal_engine::missing("strata::kernels::kv_block_bytes"); return {}; } } +namespace strata::kernels { void kv_gather_q4_step(const uint8_t * p0, const uint8_t * p1, const int32_t * p2, const int32_t * p3, const int32_t * p4, int64_t p5, const QsaShapes & p6, uint16_t * p7, uint16_t * p8, void * p9) { strata::metal_engine::missing("strata::kernels::kv_gather_q4_step"); } } +namespace strata::kernels { void kv_gather_q8_step(const int8_t * p0, const int8_t * p1, const uint16_t * p2, const uint16_t * p3, const int32_t * p4, const int32_t * p5, const int32_t * p6, int64_t p7, const QsaShapes & p8, uint16_t * p9, uint16_t * p10, void * p11) { strata::metal_engine::missing("strata::kernels::kv_gather_q8_step"); } } +namespace strata::kernels { void kv_gather_step(const uint16_t * p0, const uint16_t * p1, const int32_t * p2, const int32_t * p3, const int32_t * p4, int64_t p5, const QsaShapes & p6, uint16_t * p7, uint16_t * p8, void * p9) { strata::metal_engine::missing("strata::kernels::kv_gather_step"); } } +namespace strata::kernels { void kv_ring_restore(const QsaAttnPools & p0, const KvHostPools & p1, int p2, int64_t p3, int64_t p4, int64_t p5, const QsaShapes & p6, void * p7) { strata::metal_engine::missing("strata::kernels::kv_ring_restore"); } } +namespace strata::kernels { void kv_ring_table(int32_t * p0, int64_t p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::kernels::kv_ring_table"); } } +namespace strata::kernels { void kv_stage_from_host(const QsaAttnPools & p0, const KvHostPools & p1, int p2, int64_t p3, const QsaShapes & p4, void * p5) { strata::metal_engine::missing("strata::kernels::kv_stage_from_host"); } } +namespace strata::kernels { KvStreamCounters kv_stream_counters(const KvStreamMap & p0) { strata::metal_engine::missing("strata::kernels::kv_stream_counters"); return {}; } } +namespace strata::kernels { void kv_stream_reset(const KvStreamMap & p0, void * p1) { strata::metal_engine::missing("strata::kernels::kv_stream_reset"); } } +namespace strata::kernels { void kv_stream_resolve(const KvStreamMap & p0, const QsaAttnPools & p1, const KvHostPools & p2, int p3, const int32_t * p4, const int32_t * p5, int64_t p6, int64_t p7, const QsaShapes & p8, void * p9) { strata::metal_engine::missing("strata::kernels::kv_stream_resolve"); } } +namespace strata::kernels { void logits_nonfinite_rows(const float * p0, int p1, int p2, uint32_t * p3, void * p4) { strata::metal_engine::missing("strata::kernels::logits_nonfinite_rows"); } } +namespace strata::kernels { void map_ids(int32_t * p0, const int32_t * p1, int p2, void * p3) { strata::metal_engine::missing("strata::kernels::map_ids"); } } +namespace strata::kernels { void moe_combine(const float * p0, const float * p1, const float * p2, float * p3, int64_t p4, int64_t p5, void * p6) { strata::metal_engine::missing("strata::kernels::moe_combine"); } } +namespace strata::kernels { void moe_group_resident(const int32_t * p0, int p1, int p2, const uint8_t * p3, int64_t p4, unsigned long long * p5, int32_t * p6, int32_t * p7, int32_t * p8, int32_t * p9, void * p10) { strata::metal_engine::missing("strata::kernels::moe_group_resident"); } } +namespace strata::kernels { void moe_grouped_s2(const unsigned long long * p0, const int32_t * p1, const int32_t * p2, const int32_t * p3, const int32_t * p4, int64_t p5, int64_t p6, const uint8_t * p7, const float * p8, void * p9, float * p10, void * p11) { strata::metal_engine::missing("strata::kernels::moe_grouped_s2"); } } +namespace strata::kernels { void moe_hit_add(float * p0, const float * p1, const int32_t * p2, const int32_t * p3, int64_t p4, int64_t p5, void * p6) { strata::metal_engine::missing("strata::kernels::moe_hit_add"); } } +namespace strata::kernels { void moe_hit_grouped_s2(const uint8_t * p0, const int32_t * p1, const int32_t * p2, int64_t p3, int64_t p4, const uint8_t * p5, void * p6, float * p7, void * p8, const float * p9) { strata::metal_engine::missing("strata::kernels::moe_hit_grouped_s2"); } } +namespace strata::kernels { void moe_hit_grouped_s2_cpu_order(const uint8_t * p0, const int32_t * p1, const int32_t * p2, int64_t p3, int64_t p4, const uint8_t * p5, void * p6, float * p7, void * p8, const float * p9, float * p10) { strata::metal_engine::missing("strata::kernels::moe_hit_grouped_s2_cpu_order"); } } +namespace strata::kernels { void moe_hit_grouped_s2_dev(const uint8_t * p0, const int32_t * p1, const int32_t * p2, const int32_t * p3, int64_t p4, int64_t p5, const uint8_t * p6, void * p7, float * p8, void * p9, const float * p10) { strata::metal_engine::missing("strata::kernels::moe_hit_grouped_s2_dev"); } } +namespace strata::kernels { void moe_hit_select(const int32_t * p0, const int32_t * p1, int p2, int p3, int32_t * p4, int32_t * p5, int32_t * p6, void * p7) { strata::metal_engine::missing("strata::kernels::moe_hit_select"); } } +namespace strata::kernels { void mtp_select(const float * p0, int64_t p1, const int32_t * p2, const int32_t * p3, float * p4, int32_t * p5, int32_t * p6, int p7, void * p8, const float * p9, float * p10) { strata::metal_engine::missing("strata::kernels::mtp_select"); } } +namespace strata::kernels { bool multi_block_head_ops() { strata::metal_engine::missing("strata::kernels::multi_block_head_ops"); return {}; } } +namespace strata::kernels { void native_expert_grouped(const NativeExpertLayout & p0, const unsigned long long * p1, const int32_t * p2, const int32_t * p3, const int32_t * p4, const int32_t * p5, int64_t p6, int64_t p7, const void * p8, void * p9, float * p10, void * p11, int64_t p12) { strata::metal_engine::missing("strata::kernels::native_expert_grouped"); } } +namespace strata::kernels { void native_flash_attn_short_step(const float * p0, const uint16_t * p1, const uint16_t * p2, const int32_t * p3, int64_t p4, int p5, const QsaShapes & p6, float * p7, int32_t * p8, const uint16_t * p9, void * p10) { strata::metal_engine::missing("strata::kernels::native_flash_attn_short_step"); } } +namespace strata::kernels { void native_gdn_beta_gate(float * p0, int64_t p1, void * p2) { strata::metal_engine::missing("strata::kernels::native_gdn_beta_gate"); } } +namespace strata::kernels { void native_gdn_conv_silu(float * p0, const float * p1, const float * p2, float * p3, float * p4, int64_t p5, int64_t p6, void * p7) { strata::metal_engine::missing("strata::kernels::native_gdn_conv_silu"); } } +namespace strata::kernels { void native_gdn_gate(const float * p0, const float * p1, const float * p2, float * p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::kernels::native_gdn_gate"); } } +namespace strata::kernels { void native_gdn_l2_norm(float * p0, int64_t p1, int64_t p2, float p3, void * p4) { strata::metal_engine::missing("strata::kernels::native_gdn_l2_norm"); } } +namespace strata::kernels { void native_gdn_out_norm(const float * p0, const float * p1, const float * p2, float * p3, int64_t p4, int64_t p5, float p6, void * p7) { strata::metal_engine::missing("strata::kernels::native_gdn_out_norm"); } } +namespace strata::kernels { void native_gdn_step(float * p0, const float * p1, const float * p2, const float * p3, const float * p4, const float * p5, float * p6, const GdnShapes & p7, void * p8) { strata::metal_engine::missing("strata::kernels::native_gdn_step"); } } +namespace strata::kernels { void native_gr_post_multi(const float * p0, const float * p1, const float * p2, float * p3, int p4, int p5, int p6, long long p7, long long p8, long long p9, void * p10) { strata::metal_engine::missing("strata::kernels::native_gr_post_multi"); } } +namespace strata::kernels { void native_mmvq(int p0, const void * p1, const void * p2, float * p3, int p4, int p5, int p6, void * p7) { strata::metal_engine::missing("strata::kernels::native_mmvq"); } } +namespace strata::kernels { void native_mmvq_il(int p0, const void * p1, const void * p2, const void * p3, float * p4, int p5, int p6, int p7, void * p8) { strata::metal_engine::missing("strata::kernels::native_mmvq_il"); } } +namespace strata::kernels { void native_moe_combine(const float * p0, const float * p1, const float * p2, float * p3, int64_t p4, int64_t p5, void * p6) { strata::metal_engine::missing("strata::kernels::native_moe_combine"); } } +namespace strata::kernels { void native_moe_combine_multi(const float * p0, const float * p1, const float * p2, float * p3, int64_t p4, int64_t p5, int p6, void * p7) { strata::metal_engine::missing("strata::kernels::native_moe_combine_multi"); } } +namespace strata::kernels { void native_moe_combine_multi_hits_gated(const float * p0, const float * p1, const float * p2, const float * p3, float * p4, int64_t p5, int64_t p6, int p7, void * p8) { strata::metal_engine::missing("strata::kernels::native_moe_combine_multi_hits_gated"); } } +namespace strata::kernels { bool native_norm_rope_usable(int p0, int p1) { strata::metal_engine::missing("strata::kernels::native_norm_rope_usable"); return {}; } } +namespace strata::kernels { void native_ple_postops_batch(float * p0, float * p1, const float * p2, float * p3, const PleWeights & p4, float * p5, float * p6, float * p7, int p8, void * p9) { strata::metal_engine::missing("strata::kernels::native_ple_postops_batch"); } } +namespace strata::kernels { void native_ple_postops_batch_snap(float * p0, float * p1, const float * p2, float * p3, const PleWeights & p4, float * p5, float * p6, float * p7, int p8, float * p9, void * p10) { strata::metal_engine::missing("strata::kernels::native_ple_postops_batch_snap"); } } +namespace strata::kernels { void native_q5_k_f32(const void * p0, const float * p1, void * p2, float * p3, int p4, int p5, int p6, void * p7) { strata::metal_engine::missing("strata::kernels::native_q5_k_f32"); } } +namespace strata::kernels { bool native_q6_k_pack(const void * p0, int p1, int p2, const char * p3) { strata::metal_engine::missing("strata::kernels::native_q6_k_pack"); return {}; } } +namespace strata::kernels { void native_q8_0_pack_host(const void * p0, void * p1, int p2, int p3) { strata::metal_engine::missing("strata::kernels::native_q8_0_pack_host"); } } +namespace strata::kernels { void native_q8_0_packed_register(const void * p0, const void * p1, int p2, int p3) { strata::metal_engine::missing("strata::kernels::native_q8_0_packed_register"); } } +namespace strata::kernels { void native_q8_0_packed_unregister(const void * p0) { strata::metal_engine::missing("strata::kernels::native_q8_0_packed_unregister"); } } +namespace strata::kernels { void native_q8_1_interleave(const void * p0, void * p1, int p2, int p3, void * p4) { strata::metal_engine::missing("strata::kernels::native_q8_1_interleave"); } } +namespace strata::kernels { void native_qsa_gate_apply(const float * p0, const float * p1, float * p2, int p3, int p4, void * p5) { strata::metal_engine::missing("strata::kernels::native_qsa_gate_apply"); } } +namespace strata::kernels { void native_qsa_indexer_append(const float * p0, const int32_t * p1, int32_t p2, const float * p3, float p4, const QsaIndexerBuffers & p5, const QsaShapes & p6, int64_t p7, const RopeScaling & p8, void * p9) { strata::metal_engine::missing("strata::kernels::native_qsa_indexer_append"); } } +namespace strata::kernels { void native_qsa_indexer_append_batch(const float * p0, int64_t p1, int64_t p2, int32_t p3, const float * p4, float p5, const QsaIndexerBuffers & p6, const QsaShapes & p7, int64_t p8, const RopeScaling & p9, void * p10) { strata::metal_engine::missing("strata::kernels::native_qsa_indexer_append_batch"); } } +namespace strata::kernels { void native_qsa_indexer_append_steps(const float * p0, const int32_t * p1, int p2, int p3, int32_t p4, const float * p5, float p6, const QsaIndexerBuffers & p7, const QsaShapes & p8, int64_t p9, const RopeScaling & p10, void * p11) { strata::metal_engine::missing("strata::kernels::native_qsa_indexer_append_steps"); } } +namespace strata::kernels { void native_qsa_rms_norm_grouped(const float * p0, const float * p1, float * p2, int p3, int p4, int p5, float p6, void * p7) { strata::metal_engine::missing("strata::kernels::native_qsa_rms_norm_grouped"); } } +namespace strata::kernels { void native_qsa_rms_norm_rope(const float * p0, int p1, const float * p2, float * p3, int p4, int p5, int p6, float p7, const RopeScaling & p8, const int * p9, void * p10) { strata::metal_engine::missing("strata::kernels::native_qsa_rms_norm_rope"); } } +namespace strata::kernels { void native_qsa_rms_norm_weighted(const float * p0, const float * p1, float * p2, int p3, int p4, float p5, void * p6) { strata::metal_engine::missing("strata::kernels::native_qsa_rms_norm_weighted"); } } +namespace strata::kernels { void native_quantize_q8_1(const float * p0, void * p1, int p2, int p3, void * p4) { strata::metal_engine::missing("strata::kernels::native_quantize_q8_1"); } } +namespace strata::kernels { void native_rope_apply(const float * p0, float * p1, int p2, int p3, int p4, const RopeScaling & p5, const int * p6, void * p7) { strata::metal_engine::missing("strata::kernels::native_rope_apply"); } } +namespace strata::kernels { void native_route_resident(const float * p0, int32_t * p1, float * p2, const int32_t * p3, int p4, float p5, int p6, int p7, unsigned long long * p8, void * p9) { strata::metal_engine::missing("strata::kernels::native_route_resident"); } } +namespace strata::kernels { void native_router_top10(const float * p0, int32_t * p1, float * p2, void * p3) { strata::metal_engine::missing("strata::kernels::native_router_top10"); } } +namespace strata::kernels { void native_router_top10_multi(const float * p0, int32_t * p1, float * p2, int p3, void * p4) { strata::metal_engine::missing("strata::kernels::native_router_top10_multi"); } } +// programmatic dependent launch is a CUDA sm_90 feature: Metal never has it (hand-edited; the batch window asks this first, #3981) +namespace strata::kernels { bool & pdl_scope() { static bool s = false; return s; } } +namespace strata::kernels { bool pdl_supported() { return false; } } +namespace strata::kernels { void ple_block(const float * p0, const float * p1, const float * p2, const PleWeights & p3, PleOut & p4, void * p5, void * p6) { strata::metal_engine::missing("strata::kernels::ple_block"); } } +namespace strata::kernels { void ple_block_projected(const float * p0, const float * p1, const float * p2, const float * p3, const PleWeights & p4, PleOut & p5, void * p6, void * p7) { strata::metal_engine::missing("strata::kernels::ple_block_projected"); } } +namespace strata::kernels { void ple_history_advance(float * p0, const float * p1, void * p2) { strata::metal_engine::missing("strata::kernels::ple_history_advance"); } } +namespace strata::kernels { void qsa_attend_step(const float * p0, const uint16_t * p1, const uint16_t * p2, const int32_t * p3, int64_t p4, const QsaShapes & p5, float * p6, float * p7, void * p8) { strata::metal_engine::missing("strata::kernels::qsa_attend_step"); } } +namespace strata::kernels { void qsa_block_scores(const float * p0, const float * p1, const float * p2, const int32_t * p3, int64_t p4, int64_t p5, const QsaShapes & p6, float * p7, void * p8, int64_t p9) { strata::metal_engine::missing("strata::kernels::qsa_block_scores"); } } +namespace strata::kernels { bool qsa_block_scores_tc(const float * p0, const float * p1, const float * p2, const int32_t * p3, int64_t p4, int64_t p5, const QsaShapes & p6, float * p7, void * p8, int64_t p9) { strata::metal_engine::missing("strata::kernels::qsa_block_scores_tc"); return {}; } } +namespace strata::kernels { void qsa_block_topk(const float * p0, const int32_t * p1, int64_t p2, int64_t p3, int64_t p4, const QsaShapes & p5, int32_t * p6, void * p7, int64_t p8) { strata::metal_engine::missing("strata::kernels::qsa_block_topk"); } } +namespace strata::kernels { void qsa_decode_attn_batch(const float * p0, const QsaAttnPools & p1, const int32_t * p2, const int32_t * p3, int64_t p4, const QsaShapes & p5, float * p6, float * p7, int64_t p8, void * p9) { strata::metal_engine::missing("strata::kernels::qsa_decode_attn_batch"); } } +namespace strata::kernels { void qsa_decode_attn_step(const float * p0, const QsaAttnPools & p1, const int32_t * p2, const int32_t * p3, int64_t p4, const QsaShapes & p5, float * p6, float * p7, void * p8) { strata::metal_engine::missing("strata::kernels::qsa_decode_attn_step"); } } +namespace strata::kernels { void qsa_gate_apply_f32(const float * p0, const float * p1, const QsaShapes & p2, float * p3, void * p4) { strata::metal_engine::missing("strata::kernels::qsa_gate_apply_f32"); } } +namespace strata::kernels { void qsa_index_step(const float * p0, const float * p1, const float * p2, const QsaShapes & p3, const int32_t * p4, int64_t p5, float * p6, void * p7) { strata::metal_engine::missing("strata::kernels::qsa_index_step"); } } +namespace strata::kernels { bool qsa_prompt_attn_batch(const float * p0, const QsaAttnPools & p1, const int32_t * p2, const int32_t * p3, int64_t p4, const QsaShapes & p5, float * p6, int64_t p7, void * p8) { strata::metal_engine::missing("strata::kernels::qsa_prompt_attn_batch"); return {}; } } +namespace strata::kernels { void qsa_step_fill(int32_t * p0, int64_t p1, const QsaShapes & p2) { strata::metal_engine::missing("strata::kernels::qsa_step_fill"); } } +namespace strata::kernels { void quantize_q8_0(const float * p0, uint8_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::kernels::quantize_q8_0"); } } +namespace strata::kernels { void quantize_q8_0_scaled(const float * p0, uint8_t * p1, float * p2, int64_t p3, void * p4) { strata::metal_engine::missing("strata::kernels::quantize_q8_0_scaled"); } } +namespace strata::kernels { void quantize_q8_1_rows(const float * p0, int64_t p1, int64_t p2, void * p3, void * p4) { strata::metal_engine::missing("strata::kernels::quantize_q8_1_rows"); } } +namespace strata::kernels { void quantize_q8_K(const float * p0, uint8_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::kernels::quantize_q8_K"); } } +namespace strata::kernels { void rebase_ptrs(unsigned long long * p0, const int32_t * p1, uint8_t * p2, int64_t p3, void * p4) { strata::metal_engine::missing("strata::kernels::rebase_ptrs"); } } +namespace strata::kernels { void resident_plan(const int32_t * p0, int p1, int p2, const int32_t * p3, int p4, const uint8_t * p5, const unsigned long long * p6, long long p7, int32_t * p8, long long p9, uint32_t * p10, uint32_t p11, void * p12, uint32_t * p13) { strata::metal_engine::missing("strata::kernels::resident_plan"); } } +namespace strata::kernels { void rms_norm_weighted(float * p0, const float * p1, int64_t p2, int64_t p3, float p4, void * p5) { strata::metal_engine::missing("strata::kernels::rms_norm_weighted"); } } +namespace strata::kernels { void rope_neox_apply(const float * p0, float * p1, int64_t p2, int p3, int p4, const float * p5, const float * p6, const int * p7, void * p8) { strata::metal_engine::missing("strata::kernels::rope_neox_apply"); } } +namespace strata::kernels { void route_prior_top10(const float * p0, const int32_t * p1, float p2, int32_t * p3, float * p4, int p5, void * p6) { strata::metal_engine::missing("strata::kernels::route_prior_top10"); } } +namespace strata::kernels { void router_top10(const float * p0, int p1, int p2, int p3, int * p4, float * p5, void * p6) { strata::metal_engine::missing("strata::kernels::router_top10"); } } +namespace strata::kernels { void row_top_prob(const float * p0, int p1, int p2, const int32_t * p3, float * p4, void * p5) { strata::metal_engine::missing("strata::kernels::row_top_prob"); } } +namespace strata::kernels { void row_top_prob_split(const float * p0, int p1, int p2, const int32_t * p3, float * p4, void * p5, void * p6) { strata::metal_engine::missing("strata::kernels::row_top_prob_split"); } } +namespace strata::kernels { void s2_gemv_q8(const uint8_t * p0, const uint8_t * p1, const float * p2, float * p3, int64_t p4, int64_t p5, int p6, void * p7) { strata::metal_engine::missing("strata::kernels::s2_gemv_q8"); } } +namespace strata::kernels { void s_gemv_q8_0_split(const uint8_t * p0, const uint8_t * p1, const float * p2, const float * p3, float * p4, int64_t p5, int64_t p6, const SForm & p7, void * p8) { strata::metal_engine::missing("strata::kernels::s_gemv_q8_0_split"); } } +namespace strata::kernels { void s_gemv_q8k_split(const uint8_t * p0, const uint8_t * p1, const float * p2, const float * p3, float * p4, int64_t p5, int64_t p6, const SForm & p7, void * p8) { strata::metal_engine::missing("strata::kernels::s_gemv_q8k_split"); } } +namespace strata::kernels { void sample_tokens(const float * p0, int p1, int p2, const int * p3, int p4, const SamplerParams & p5, int * p6, void * p7) { strata::metal_engine::missing("strata::kernels::sample_tokens"); } } +namespace strata::kernels { void scale_inplace(float * p0, int64_t p1, float p2, void * p3) { strata::metal_engine::missing("strata::kernels::scale_inplace"); } } +namespace strata::kernels { void scatter_rows_f32(const float * p0, float * p1, const int32_t * p2, int64_t p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::kernels::scatter_rows_f32"); } } +namespace strata::kernels { void shared_expert(const uint8_t * p0, const uint8_t * p1, const uint16_t * p2, const SForm & p3, const uint8_t * p4, const float * p5, const float * p6, const SForm & p7, const uint8_t * p8, const float * p9, const float * p10, const SForm & p11, const uint8_t * p12, const float * p13, const float * p14, const uint16_t * p15, float * p16, float * p17, int64_t p18, int64_t p19, int p20, void * p21, const float * p22, const NativeSharedWeights * p23) { strata::metal_engine::missing("strata::kernels::shared_expert"); } } +namespace strata::kernels { void shared_expert_multi(int p0, const float * p1, const uint16_t * p2, const NativeSharedWeights & p3, const uint16_t * p4, float * p5, float * p6, float * p7, float * p8, int64_t p9, int64_t p10, void * p11, const void * p12, int p13) { strata::metal_engine::missing("strata::kernels::shared_expert_multi"); } } +namespace strata::kernels { void silu_inplace(float * p0, int64_t p1, void * p2) { strata::metal_engine::missing("strata::kernels::silu_inplace"); } } +namespace strata::kernels { void spec_draft_sample(float * p0, int p1, const int32_t * p2, const int32_t * p3, int p4, const SamplerParams * p5, int32_t * p6, int p7, int p8, const int32_t * p9, void * p10, int32_t * p11, float * p12, int32_t * p13, int p14, float p15, void * p16) { strata::metal_engine::missing("strata::kernels::spec_draft_sample"); } } +namespace strata::kernels { void topk_512_step(const float * p0, const QsaShapes & p1, int64_t p2, const int32_t * p3, int32_t * p4, void * p5) { strata::metal_engine::missing("strata::kernels::topk_512_step"); } } +namespace strata::kernels { void wait_flag_ge(const uint32_t * p0, uint32_t p1, void * p2) { strata::metal_engine::missing("strata::kernels::wait_flag_ge"); } } +namespace strata::kernels { void wait_flag_ge_or(const uint32_t * p0, uint32_t p1, const uint32_t * p2, void * p3) { strata::metal_engine::missing("strata::kernels::wait_flag_ge_or"); } } +namespace strata::kernels { void window_ids(int32_t * p0, int p1, int p2, int32_t * p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::kernels::window_ids"); } } +namespace strata::kernels { void window_ids_floor(int32_t * p0, int p1, int p2, int32_t * p3, int64_t p4, const int32_t * p5, void * p6) { strata::metal_engine::missing("strata::kernels::window_ids_floor"); } } +auto strata::prefill::Gemm::bf16(const uint16_t * p0, const uint16_t * p1, float * p2, int64_t p3, int64_t p4, int64_t p5, int64_t p6, float p7, int64_t p8, bool p9) -> void { strata::metal_engine::missing("strata::prefill::Gemm::bf16"); } +auto strata::prefill::Gemm::bf16_converts_x() const -> bool { strata::metal_engine::missing("strata::prefill::Gemm::bf16_converts_x"); return {}; } +auto strata::prefill::Gemm::bf16_hcd_exact(const uint16_t * p0, int64_t p1, const uint16_t * p2, float * p3, int64_t p4, int64_t p5, int64_t p6) -> bool { strata::metal_engine::missing("strata::prefill::Gemm::bf16_hcd_exact"); return {}; } +auto strata::prefill::Gemm::f16(const uint16_t * p0, const uint16_t * p1, float * p2, int64_t p3, int64_t p4, int64_t p5, int64_t p6, float p7) -> void { strata::metal_engine::missing("strata::prefill::Gemm::f16"); } +auto strata::prefill::Gemm::native(const uint16_t * p0, int p1, const void * p2, float * p3, int64_t p4, int64_t p5, int64_t p6, int64_t p7, float p8, int64_t p9) -> void { strata::metal_engine::missing("strata::prefill::Gemm::native"); } +auto strata::prefill::Gemm::rebind(uint16_t * p0, int64_t p1, void * p2, size_t p3) -> void { strata::metal_engine::missing("strata::prefill::Gemm::rebind"); } +namespace strata::prefill { void blob_dequant_f16(const uint8_t * p0, uint16_t * p1, uint16_t * p2, void * p3) { strata::metal_engine::missing("strata::prefill::blob_dequant_f16"); } } +namespace strata::prefill { void copy_f32_wide(float * p0, const float * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::prefill::copy_f32_wide"); } } +namespace strata::prefill { void copy_i32(int32_t * p0, const int32_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::prefill::copy_i32"); } } +namespace strata::prefill { void f16_to_f32_wide(float * p0, const uint16_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::prefill::f16_to_f32_wide"); } } +namespace strata::prefill { void gate_attn(const float * p0, const float * p1, uint16_t * p2, int64_t p3, void * p4, int64_t p5) { strata::metal_engine::missing("strata::prefill::gate_attn"); } } +namespace strata::prefill { void gather_rows16(const uint16_t * p0, const int32_t * p1, uint16_t * p2, int64_t p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::prefill::gather_rows16"); } } +namespace strata::prefill { void gdn_conv(float * p0, const float * p1, const float * p2, float * p3, int64_t p4, float p5, void * p6) { strata::metal_engine::missing("strata::prefill::gdn_conv"); } } +namespace strata::prefill { void gdn_gates(const float * p0, const float * p1, const float * p2, float * p3, float * p4, int64_t p5, void * p6) { strata::metal_engine::missing("strata::prefill::gdn_gates"); } } +namespace strata::prefill { void gdn_recurrence(float * p0, const float * p1, const float * p2, const float * p3, const float * p4, const float * p5, float p6, float * p7, uint16_t * p8, int64_t p9, void * p10, int64_t p11) { strata::metal_engine::missing("strata::prefill::gdn_recurrence"); } } +namespace strata::prefill { void gdn_recurrence_state_only(float * p0, const float * p1, const float * p2, const float * p3, float * p4, int64_t p5, void * p6) { strata::metal_engine::missing("strata::prefill::gdn_recurrence_state_only"); } } +namespace strata::prefill { void gr_broadcast(const float * p0, float * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::prefill::gr_broadcast"); } } +namespace strata::prefill { void gr_mix(const float * p0, const float * p1, float * p2, uint16_t * p3, int64_t p4, void * p5, uint16_t * p6, uint16_t * p7) { strata::metal_engine::missing("strata::prefill::gr_mix"); } } +namespace strata::prefill { void gr_mix_r(const float * p0, const float * p1, const float * p2, const float * p3, float * p4, uint16_t * p5, int64_t p6, void * p7, uint16_t * p8, uint16_t * p9) { strata::metal_engine::missing("strata::prefill::gr_mix_r"); } } +namespace strata::prefill { void gr_norm(const float * p0, const float * p1, float p2, float * p3, uint16_t * p4, int64_t p5, void * p6, uint16_t * p7) { strata::metal_engine::missing("strata::prefill::gr_norm"); } } +namespace strata::prefill { void gr_norm_rs(const float * p0, const float * p1, float p2, float * p3, uint16_t * p4, int64_t p5, void * p6, uint16_t * p7, int64_t p8) { strata::metal_engine::missing("strata::prefill::gr_norm_rs"); } } +namespace strata::prefill { void gr_silu(const float * p0, uint16_t * p1, int64_t p2, void * p3, uint16_t * p4) { strata::metal_engine::missing("strata::prefill::gr_silu"); } } +namespace strata::prefill { bool gr_upmix(const uint16_t * p0, const uint16_t * p1, const float * p2, const float * p3, const float * p4, float * p5, uint16_t * p6, uint16_t * p7, int64_t p8, void * p9) { strata::metal_engine::missing("strata::prefill::gr_upmix"); return {}; } } +namespace strata::prefill { void gr_write(float * p0, const float * p1, const float * p2, int64_t p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::prefill::gr_write"); } } +namespace strata::prefill { void gr_write_cvec_norm_rs(float * p0, const float * p1, const float * p2, int64_t p3, const float * p4, const float * p5, const int * p6, int p7, const float * p8, float p9, float * p10, uint16_t * p11, int64_t p12, void * p13, uint16_t * p14, int64_t p15) { strata::metal_engine::missing("strata::prefill::gr_write_cvec_norm_rs"); } } +namespace strata::prefill { void gr_write_norm_rs(float * p0, const float * p1, const float * p2, int64_t p3, const float * p4, float p5, float * p6, uint16_t * p7, int64_t p8, void * p9, uint16_t * p10, int64_t p11) { strata::metal_engine::missing("strata::prefill::gr_write_norm_rs"); } } +namespace strata::prefill { void kv_append(const float * p0, const float * p1, int64_t p2, int64_t p3, const int32_t * p4, int64_t p5, uint16_t * p6, uint16_t * p7, int8_t * p8, int8_t * p9, uint16_t * p10, uint16_t * p11, void * p12, const strata::kernels::KvHostPools * p13, const strata::kernels::KvHostPools * p14) { strata::metal_engine::missing("strata::prefill::kv_append"); } } +namespace strata::prefill { void moe_combine(const float * p0, const int32_t * p1, const float * p2, const float * p3, const float * p4, float * p5, int64_t p6, void * p7) { strata::metal_engine::missing("strata::prefill::moe_combine"); } } +namespace strata::prefill { void moe_combine_peer(const float * p0, const int32_t * p1, const float * p2, const float * p3, const float * p4, const float * p5, int64_t p6, float * p7, int64_t p8, void * p9) { strata::metal_engine::missing("strata::prefill::moe_combine_peer"); } } +namespace strata::prefill { void moe_combine_peer16(const float * p0, const int32_t * p1, const float * p2, const float * p3, const float * p4, const uint16_t * p5, int64_t p6, float * p7, int64_t p8, void * p9) { strata::metal_engine::missing("strata::prefill::moe_combine_peer16"); } } +namespace strata::prefill { void peer_gather_add(float * p0, const float * p1, int64_t p2, const float * p3, const int32_t * p4, const int32_t * p5, const int32_t * p6, const int32_t * p7, int64_t p8, void * p9) { strata::metal_engine::missing("strata::prefill::peer_gather_add"); } } +namespace strata::prefill { void peer_scatter_add(float * p0, const float * p1, const float * p2, const int32_t * p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::prefill::peer_scatter_add"); } } +namespace strata::prefill { void rms_rows(float * p0, const float * p1, int64_t p2, int64_t p3, int64_t p4, float p5, void * p6) { strata::metal_engine::missing("strata::prefill::rms_rows"); } } +namespace strata::prefill { void rope(float * p0, int64_t p1, int64_t p2, int64_t p3, int64_t p4, int64_t p5, const strata::kernels::RopeScaling & p6, void * p7) { strata::metal_engine::missing("strata::prefill::rope"); } } +namespace strata::prefill { void round_f16(const float * p0, float * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::prefill::round_f16"); } } +namespace strata::prefill { void route(const float * p0, int32_t * p1, float * p2, int64_t p3, int64_t p4, void * p5) { strata::metal_engine::missing("strata::prefill::route"); } } +namespace strata::prefill { void split_q(const float * p0, float * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::prefill::split_q"); } } +namespace strata::prefill { void sums_to_f16(const float * p0, uint16_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::prefill::sums_to_f16"); } } +namespace strata::prefill { void swiglu_interleaved(const float * p0, uint16_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::prefill::swiglu_interleaved"); } } +namespace strata::prefill { void swiglu_pair(const float * p0, const float * p1, uint16_t * p2, int64_t p3, void * p4) { strata::metal_engine::missing("strata::prefill::swiglu_pair"); } } +namespace strata::prefill { void to_bf16(const float * p0, uint16_t * p1, int64_t p2, void * p3, uint16_t * p4) { strata::metal_engine::missing("strata::prefill::to_bf16"); } } +namespace strata::prefill { void to_f16(const float * p0, uint16_t * p1, int64_t p2, void * p3) { strata::metal_engine::missing("strata::prefill::to_f16"); } } diff --git a/edge0-strata/metal/src/engine/stubs_0145.cpp b/edge0-strata/metal/src/engine/stubs_0145.cpp new file mode 100644 index 00000000..722459f4 --- /dev/null +++ b/edge0-strata/metal/src/engine/stubs_0145.cpp @@ -0,0 +1,165 @@ +// metal/src/engine/stubs_0145.cpp - the kernel-library entry points the 0.1.45 sources added and Metal does not implement yet. +// Same job and shape as stubs.cpp (each stops the engine with its name, metal/src/engine/missing.cpp). Written from the linker's +// undefined-symbol report because tools/gen_stubs.py reads an older clang AST dump format than macOS 27's clang prints. +// A symbol that gets a Metal implementation is deleted from here. fused_gr_persist_init / fused_gr_read_persist / f16_range_read / +// gemm_f16_range_read are opt-in features of the HIP and Turing builds: off on Metal, so they answer "nothing to do" instead of stopping the engine. + +#include "strata/kernels/bf16_bits.hpp" +#include "strata/kernels/bf16_gemv.hpp" +#include "strata/kernels/cvec.hpp" +#include "strata/kernels/dequant_bf16.hpp" +#include "strata/kernels/dequant_s2.hpp" +#include "strata/kernels/elementwise.hpp" +#include "strata/kernels/f16_bits.hpp" +#include "strata/kernels/fused_gdn.hpp" +#include "strata/kernels/fused_gr.hpp" +#include "strata/kernels/gdn.hpp" +#include "strata/kernels/gfx_arch.hpp" +#include "strata/kernels/gr.hpp" +#include "strata/kernels/iq_kernels.hpp" +#include "strata/kernels/kv_q4.hpp" +#include "strata/kernels/kv_q8.hpp" +#include "strata/kernels/kv_stream.hpp" +#include "strata/kernels/mrope.hpp" +#include "strata/kernels/native_flash_attn.hpp" +#include "strata/kernels/native_gdn.hpp" +#include "strata/kernels/native_gdn_preprocess.hpp" +#include "strata/kernels/native_gr_norm.hpp" +#include "strata/kernels/native_gr_postops.hpp" +#include "strata/kernels/native_mmvq.hpp" +#include "strata/kernels/native_moe.hpp" +#include "strata/kernels/native_ple_postops.hpp" +#include "strata/kernels/native_qsa.hpp" +#include "strata/kernels/native_qsa_indexer.hpp" +#include "strata/kernels/native_qsa_score.hpp" +#include "strata/kernels/native_rope.hpp" +#include "strata/kernels/native_router.hpp" +#include "strata/kernels/ngram.hpp" +#include "strata/kernels/pdl.hpp" +#include "strata/kernels/ple.hpp" +#include "strata/kernels/qsa.hpp" +#include "strata/kernels/qsa_decode_attn.hpp" +#include "strata/kernels/qsa_prompt_attn.hpp" +#include "strata/kernels/qsa_select.hpp" +#include "strata/kernels/quantize_act.hpp" +#include "strata/kernels/rope.hpp" +#include "strata/kernels/rope_scaling.hpp" +#include "strata/kernels/route_prior.hpp" +#include "strata/kernels/router_top10.hpp" +#include "strata/kernels/s2_expert_grouped.hpp" +#include "strata/kernels/s2_gemv.hpp" +#include "strata/kernels/s2_gemv_q8.hpp" +#include "strata/kernels/s_gemv.hpp" +#include "strata/kernels/sampler.hpp" +#include "strata/kernels/shared_expert.hpp" +#include "strata/kernels/verify_kernels.hpp" +#include "strata/kernels/cpu/expert.hpp" +#include "strata/kernels/cpu/expert_layout.hpp" +#include "strata/kernels/cpu/iq_avx1.hpp" +#include "strata/kernels/cpu/iq_avx2.hpp" +#include "strata/kernels/cpu/iq_avx512.hpp" +#include "strata/kernels/cpu/kq_avx1.hpp" +#include "strata/kernels/cpu/kq_avx2.hpp" +#include "strata/kernels/cpu/native_expert.hpp" +#include "strata/kernels/cpu/pool.hpp" +#include "strata/prefill/gemm.hpp" +#include "strata/prefill/kernels.hpp" +#include "strata/prefill/moe_fused.hpp" +#include "strata/prefill/moe_fused_iq.hpp" +#include "strata/prefill/moe_mmq.hpp" +#include "strata/prefill/prefill.hpp" +#include "strata/prefill/share_rules.hpp" +#include "strata/prefill/xmx_moe.hpp" +#include "strata/core/arch_defaults.hpp" +#include "strata/core/batch_rows.hpp" +#include "strata/core/conversation_buffer.hpp" +#include "strata/core/conversation_cache.hpp" +#include "strata/core/conversation_file.hpp" +#include "strata/core/conversation_memory.hpp" +#include "strata/core/conversation_snapshot.hpp" +#include "strata/core/coupled_draft.hpp" +#include "strata/core/device.hpp" +#include "strata/core/dma_batch.hpp" +#include "strata/core/emulate.hpp" +#include "strata/core/exchange_storage.hpp" +#include "strata/core/expert_cache.hpp" +#include "strata/core/expert_source.hpp" +#include "strata/core/foresight_swap.hpp" +#include "strata/core/graph.hpp" +#include "strata/core/hit_hook.hpp" +#include "strata/core/layer.hpp" +#include "strata/core/layout.hpp" +#include "strata/core/logit_bias.hpp" +#include "strata/core/mtp.hpp" +#include "strata/core/native_dense.hpp" +#include "strata/core/native_head.hpp" +#include "strata/core/on_device.hpp" +#include "strata/core/peer_experts.hpp" +#include "strata/core/pinned.hpp" +#include "strata/core/pinned_ring.hpp" +#include "strata/core/progress.hpp" +#include "strata/core/remote_expert_opt.hpp" +#include "strata/core/remote_experts.hpp" +#include "strata/core/session.hpp" +#include "strata/core/session_save_reclaim.hpp" +#include "strata/core/spec_prob.hpp" +#include "strata/core/verify.hpp" +#include "strata/core/vmm.hpp" +#include "strata/core/weights.hpp" +#include "strata/platform/aux_cpus.hpp" +#include "strata/platform/direct_file.hpp" +#include "strata/platform/host_limits.hpp" +#include "strata/platform/memory.hpp" +#include "strata/ngram/ple_reader.hpp" +#include "strata/spec/controller.hpp" +#include "strata/spec/draft_policy.hpp" +#include "strata/spec/draft_source.hpp" +#include "strata/spec/suffix_drafter.hpp" +#include + +#include + +namespace strata::metal_engine { void missing(const char* name); } + +namespace strata::kernels { void copy_rows_from_mapped_after(unsigned int const* p0, unsigned int p1, float* p2, float const* p3, long long p4, long long p5, int const* p6, int const* p7, void* p8) { strata::metal_engine::missing("strata::kernels::copy_rows_from_mapped_after"); } } +namespace strata::kernels::cpu { void iq256_gu_rows_exact_one(int p0, unsigned char const* p1, unsigned long p2, unsigned long p3, int p4, void const* p5, float* p6, int p7, int p8) { strata::metal_engine::missing("strata::kernels::cpu::iq256_gu_rows_exact_one"); } } +namespace strata::kernels { void fused_gr_persist_init() { } } +namespace strata::kernels { bool fused_gr_read_persist(strata::kernels::FusedGrArgs const* p0, int p1, void* p2) { return false; } } +namespace strata::kernels { void gdn_step_norm_groups(strata::kernels::GdnGroups const& p0, float const* p1, int p2, float const* p3, float const* p4, float const* p5, float const* p6, float p7, float* p8, int p9, int p10, bool p11, void* p12) { strata::metal_engine::missing("strata::kernels::gdn_step_norm_groups"); } } +namespace strata::kernels { void handoff_publish(float* p0, float const* p1, long long p2, float const* p3, long long p4, float const* p5, long long p6, unsigned int* p7, unsigned int* p8, void* p9, unsigned int const* p10) { strata::metal_engine::missing("strata::kernels::handoff_publish"); } } +namespace strata::kernels { void hc_q8_to_bf16(void const* p0, float const* p1, unsigned short* p2, long long p3, void* p4) { strata::metal_engine::missing("strata::kernels::hc_q8_to_bf16"); } } +namespace strata::kernels { void native_quantize_q8_1_il(float const* p0, void* p1, void* p2, int p3, int p4, void* p5) { strata::metal_engine::missing("strata::kernels::native_quantize_q8_1_il"); } } +namespace strata::kernels { void qsa_block_scores_rows(strata::kernels::QsaScoreRows const& p0, float const* p1, int const* p2, long long p3, long long p4, strata::kernels::QsaShapes const& p5, float* p6, void* p7) { strata::metal_engine::missing("strata::kernels::qsa_block_scores_rows"); } } +namespace strata::kernels { void qsa_decode_attn_batch_rows(float const* p0, strata::kernels::QsaAttnPoolsRows const& p1, int const* p2, int const* p3, long long p4, strata::kernels::QsaShapes const& p5, float* p6, float* p7, long long p8, void* p9) { strata::metal_engine::missing("strata::kernels::qsa_decode_attn_batch_rows"); } } +namespace strata::kernels { bool qsa_decode_attn_split(float const* p0, strata::kernels::QsaAttnPools const& p1, int const* p2, int const* p3, long long p4, strata::kernels::QsaShapes const& p5, float* p6, unsigned long long p7, float* p8, long long p9, void* p10) { strata::metal_engine::missing("strata::kernels::qsa_decode_attn_split"); return {}; } } +namespace strata::kernels { void res_patch(int* p0, int const* p1, int p2, void* p3) { strata::metal_engine::missing("strata::kernels::res_patch"); } } +namespace strata::kernels { void wait_copy_i32_from_mapped(unsigned int const* p0, unsigned int p1, int* p2, int const* p3, long long p4, void* p5) { strata::metal_engine::missing("strata::kernels::wait_copy_i32_from_mapped"); } } +namespace strata::kernels { void wait_flag_ge_stamped(unsigned int const* p0, unsigned int p1, unsigned long long* p2, unsigned long long* p3, void* p4) { strata::metal_engine::missing("strata::kernels::wait_flag_ge_stamped"); } } +namespace strata::kernels { void wait_handoff(unsigned int const* p0, unsigned int p1, unsigned int* p2, unsigned long long p3, void* p4) { strata::metal_engine::missing("strata::kernels::wait_handoff"); } } +namespace strata::prefill { void moe_scatter_add(float const* p0, int const* p1, int const* p2, int const* p3, float const* p4, long long p5, float* p6, void* p7) { strata::metal_engine::missing("strata::prefill::moe_scatter_add"); } } +namespace strata::prefill { void moe_shared_init(float const* p0, float const* p1, float* p2, long long p3, void* p4) { strata::metal_engine::missing("strata::prefill::moe_shared_init"); } } +namespace strata::prefill { bool f16_range_read(float& p0, unsigned long long& p1, bool p2) { return false; } } +namespace strata::prefill { bool gemm_f16_range_read(float* p0, unsigned long long* p1, bool p2) { return false; } } +namespace strata::prefill { void gr_norm(float const* p0, float const* p1, float p2, float* p3, unsigned short* p4, long long p5, void* p6, unsigned short* p7, bool p8) { strata::metal_engine::missing("strata::prefill::gr_norm"); } } +namespace strata::prefill { void gr_norm_rs(float const* p0, float const* p1, float p2, float* p3, unsigned short* p4, long long p5, void* p6, unsigned short* p7, long long p8, bool p9) { strata::metal_engine::missing("strata::prefill::gr_norm_rs"); } } +namespace strata::prefill { void gr_write_cvec_norm_rs(float* p0, float const* p1, float const* p2, long long p3, float const* p4, float const* p5, int const* p6, int p7, float const* p8, float p9, float* p10, unsigned short* p11, long long p12, void* p13, unsigned short* p14, long long p15, bool p16) { strata::metal_engine::missing("strata::prefill::gr_write_cvec_norm_rs"); } } +namespace strata::prefill { void gr_write_norm_rs(float* p0, float const* p1, float const* p2, long long p3, float const* p4, float p5, float* p6, unsigned short* p7, long long p8, void* p9, unsigned short* p10, long long p11, bool p12) { strata::metal_engine::missing("strata::prefill::gr_write_norm_rs"); } } +// pf_wmma (STRATA_PF_WMMA=1, Volta prompt experts) is CUDA only; the prefill code of 148/149 references it, so the Metal link needs the entry point. Off on Metal: "not launched". +namespace strata::kernels { bool gemm_iq_f16_grouped(int p0, void const* p1, unsigned long p2, int p3, int p4, void const* p5, int p6, int const* p7, int p8, int p9, float* p10, int p11, void* p12) { return false; } } +// STRATA_Q8_SM60 (Pascal GP100, opt-in): a no-op on Metal; NativeHead::load()/unload() call these for a Q8_0 head, so they must exist and answer "nothing packed". +namespace strata::kernels { bool native_q8_0_sm60_pack(void const* p0, int p1, int p2, char const* p3) { return false; } } +namespace strata::kernels { void native_q8_0_sm60_release(void const* p0) { } } +// the MTP drafter's STRATA_MTP_Q8_SM60 packing (mtp.cpp): nothing packed on Metal +namespace strata::kernels { bool native_q8_0_sm60_pack_quiet(void const* p0, int p1, int p2, char const* p3) { return false; } } +namespace strata::kernels { bool native_q8_0_sm60_pack_in_place(void* p0, int p1, int p2, char const* p3) { return false; } } +// STRATA_SELECT_LONG (verify.cpp's long-context decode select): CUDA only; never reached on Metal (no scratch asked for) +namespace strata::kernels { long long qsa_topk_multi_scratch_ints(long long p0) { return 0; } } +namespace strata::kernels { void qsa_block_topk_long(float const* p0, int const* p1, long long p2, long long p3, long long p4, strata::kernels::QsaShapes const& p5, int* p6, void* p7, long long p8, int* p9) { strata::metal_engine::missing("strata::kernels::qsa_block_topk_long"); } } +namespace strata::kernels { void qsa_block_scores_long(float const* p0, float const* p1, float const* p2, int const* p3, long long p4, long long p5, strata::kernels::QsaShapes const& p6, float* p7, void* p8, long long p9) { strata::metal_engine::missing("strata::kernels::qsa_block_scores_long"); } } +// only reached with pf_wmma (off on Metal) +namespace strata::prefill { void swiglu_split_f16(float const* p0, unsigned short* p1, long long p2, int p3, bool p4, void* p5) { } } +// STRATA_PLE_COMPACT (opt-in, CUDA / HIP prompt path): prefill.cpp references it; the Metal prompt runs through its window and never reaches it. +namespace strata::kernels { void ple_iq4nl_dequant(const uint8_t* p0, int64_t p1, float* p2, void* p3) { strata::metal_engine::missing("strata::kernels::ple_iq4nl_dequant"); } } +// STRATA_DEQUANT_SIDE (opt-in, CUDA only): off on Metal; the prompt path takes no second scratch and set_dequant_side(nullptr) keeps the serial form. +namespace strata::prefill { bool gemm_dequant_side() { return false; } } +auto strata::prefill::Gemm::set_dequant_side(uint16_t* p0, std::string& p1) -> bool { if (p0 == nullptr) return true; strata::metal_engine::missing("strata::prefill::Gemm::set_dequant_side"); return false; } diff --git a/edge0-strata/metal/src/engine/stubs_023.cpp b/edge0-strata/metal/src/engine/stubs_023.cpp new file mode 100644 index 00000000..49aa1692 --- /dev/null +++ b/edge0-strata/metal/src/engine/stubs_023.cpp @@ -0,0 +1,15 @@ +// metal/src/engine/stubs_023.cpp - the two research hooks of the shared Verifier::record_window that only the CUDA / HIP backends implement +// (src/kernels/cuda/verify_kernels.cu): hc_mean_streams (STRATA_TAP_FILE) and route_marginal_stats (STRATA_ROUTE_STATS). Verifier::record_window +// does not run on Metal (init hands over to the Metal window, metal/src/engine/window.cpp), so neither is reached; if one ever is, it stops +// the engine with its name instead of continuing on outputs nobody wrote. They exist so the shared host file links. + +#include + +#include "strata/kernels/verify_kernels.hpp" + +namespace strata::metal_engine { void missing(const char* name); } + +namespace strata::kernels { +void hc_mean_streams(const float*, float*, int64_t, int, void*) { strata::metal_engine::missing("strata::kernels::hc_mean_streams"); } +void route_marginal_stats(const int32_t*, int, int, unsigned long long*, void*) { strata::metal_engine::missing("strata::kernels::route_marginal_stats"); } +} // namespace strata::kernels diff --git a/edge0-strata/metal/src/engine/stubs_025.cpp b/edge0-strata/metal/src/engine/stubs_025.cpp new file mode 100644 index 00000000..6e9f3bfe --- /dev/null +++ b/edge0-strata/metal/src/engine/stubs_025.cpp @@ -0,0 +1,19 @@ +// metal/src/engine/stubs_025.cpp - 0.2.5: native_expert.cpp (shared with the PC engine) now asks whether the 512-bit i-quant rows may take a +// one-token expert and whether their grid lookups are gathers. Those rows live in src/kernels/cpu/iq_avx512.cpp, an x86 AVX-512 file the +// Apple build does not compile. Here they answer "not available"; native_expert.cpp only calls iq512x_rows behind cpu_avx512_ok() and +// iq512x_supported(), both false on arm64, so the stopping stub below is never reached. + +#include +#include + +namespace strata::metal_engine { void missing(const char* name); } + +namespace strata::kernels::cpu { + +bool iq512_gather_on() noexcept { return false; } +bool iq512x_supported(int) noexcept { return false; } +void iq512x_rows(int, const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int) { + strata::metal_engine::missing("strata::kernels::cpu::iq512x_rows"); +} + +} // namespace strata::kernels::cpu diff --git a/edge0-strata/metal/src/engine/stubs_144.cpp b/edge0-strata/metal/src/engine/stubs_144.cpp new file mode 100644 index 00000000..a4a151b7 --- /dev/null +++ b/edge0-strata/metal/src/engine/stubs_144.cpp @@ -0,0 +1,64 @@ +// metal/src/engine/stubs_144.cpp - the 0.1.44 PC entry points the Metal engine does not implement (merge of the Mac branch +// with the PC release): new overloads / new functions that stubs.cpp (generated) does not know. Each either stops the engine +// with its name (like stubs.cpp) or, where the caller has a documented fallback, answers "not available". None is reached +// with the Metal defaults: STRATA_COMMIT_BATCH is off on Metal, the NVFP4 CPU rows need AVX2 / AVX-512, and the K-quant / +// NVFP4 prompt paths are CUDA / HIP only. + +#include +#include + +#include "strata/kernels/cpu/nvfp4_avx2.hpp" +#include "strata/kernels/cpu/nvfp4_avx512.hpp" +#include "strata/kernels/iq_kernels.hpp" +#include "strata/kernels/native_qsa_indexer.hpp" +#include "strata/kernels/qsa_decode_attn.hpp" +#include "strata/kernels/qsa_select.hpp" +#include "strata/kernels/verify_kernels.hpp" +#include "strata/prefill/kernels.hpp" + +namespace strata::metal_engine { +void missing(const char* name); +} + +namespace strata::kernels { +void iq_dequant_f16(int, const void*, int64_t, uint16_t*, void*, const float*) { strata::metal_engine::missing("strata::kernels::iq_dequant_f16(scale)"); } +void iq_dequant_gu_f16(int, const void*, const void*, int64_t, int64_t, uint16_t*, void*, const float*) { strata::metal_engine::missing("strata::kernels::iq_dequant_gu_f16(tail)"); } +bool gdn_commit_layers(float*, int64_t, int64_t, const float*, int64_t, const float*, int64_t, const float*, const float*, int64_t, + int, int, int, int, int, const int32_t*, void*) { return false; } // "nothing launched: keep the per-layer loop" +void native_qsa_indexer_append_steps_layers(const float*, int64_t, const int32_t*, int, int, int32_t, float, const QsaIndexerLayers&, int, + const QsaShapes&, int64_t, const RopeScaling&, void*) { + strata::metal_engine::missing("strata::kernels::native_qsa_indexer_append_steps_layers"); +} +} // namespace strata::kernels + +namespace strata::prefill { +void moe_combine(const float*, const int32_t*, const float*, const float*, const float*, float*, int64_t, void*, const float*) { + strata::metal_engine::missing("strata::prefill::moe_combine(row_sd)"); +} +} // namespace strata::prefill + +namespace strata::kernels::cpu { +bool nvfp4_256_fits(int) { return false; } +bool nvfp4_512_fits(int) { return false; } +void nvfp4_256_gu_rows(const uint8_t*, size_t, size_t, int, const void* const*, int, float* const*, int, int, float, float) { + strata::metal_engine::missing("strata::kernels::cpu::nvfp4_256_gu_rows"); +} +void nvfp4_256_rows(const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int, float) { + strata::metal_engine::missing("strata::kernels::cpu::nvfp4_256_rows"); +} +void nvfp4_512_gu_rows(const uint8_t*, size_t, size_t, int, const void* const*, int, float* const*, int, int, float, float) { + strata::metal_engine::missing("strata::kernels::cpu::nvfp4_512_gu_rows"); +} +void nvfp4_512_rows(const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int, float) { + strata::metal_engine::missing("strata::kernels::cpu::nvfp4_512_rows"); +} +} // namespace strata::kernels::cpu + +// STRATA_QSA_SLOTS (PR #2364): the one-launch-over-the-rows QSA launchers have no Metal kernel; false = "run the per-row calls", +// which is what Verifier::record_window does on a decline. +namespace strata::kernels { +bool qsa_block_scores_slots(const float* const*, const float* const*, const float*, const int32_t*, int64_t, int64_t, + const QsaShapes&, float*, void*) { return false; } +bool qsa_decode_attn_slots(const float*, const QsaAttnSlotPools&, const int32_t*, const int32_t*, int64_t, const QsaShapes&, + float*, float*, int64_t, void*) { return false; } +} // namespace strata::kernels diff --git a/edge0-strata/metal/src/engine/stubs_ds4.cpp b/edge0-strata/metal/src/engine/stubs_ds4.cpp new file mode 100644 index 00000000..9abafef0 --- /dev/null +++ b/edge0-strata/metal/src/engine/stubs_ds4.cpp @@ -0,0 +1,16 @@ +// metal/src/engine/stubs_ds4.cpp - CPU-kernel entry points the DeepSeek engine (strata_ds4) references that have no Apple +// implementation. strata::ds4::gguf::dots() asks iq256x_supported() before it uses iq256x_rows(); on arm64 the answer is no and the +// CPU lane keeps ggml-cpu's rows, which have their own NEON kernels. +#include "strata/kernels/cpu/iq_avx2.hpp" + +#include + +namespace strata::kernels::cpu { + +bool iq256x_supported(int) noexcept { return false; } + +void iq256x_rows(int, const uint8_t*, size_t, int, const void* const*, int, float* const*, int, int) { + std::abort(); // unreachable: iq256x_supported() is false here, and callers check it +} + +} // namespace strata::kernels::cpu diff --git a/edge0-strata/metal/src/engine/stubs_gyro.cpp b/edge0-strata/metal/src/engine/stubs_gyro.cpp new file mode 100644 index 00000000..4c2f336e --- /dev/null +++ b/edge0-strata/metal/src/engine/stubs_gyro.cpp @@ -0,0 +1,25 @@ +// metal/src/engine/stubs_gyro.cpp - the Gyro (agentionai trellis experts) entry points the Metal engine does not implement: +// the GPU Hadamard rotations of src/kernels/cuda/iq_kernels.cu. native_expert_supported() (iq_host.cpp) does not list the +// trellis types, so a Gyro pack stops at load with its reason and these are never reached with a supported file. +// Each stops the engine with its name, like stubs.cpp. + +#include + +#include "strata/kernels/iq_kernels.hpp" + +namespace strata::metal_engine { +void missing(const char* name); +} + +namespace strata::kernels { +void hadamard_rows(const float*, float*, int64_t, int64_t, int, const float*, void*) { + strata::metal_engine::missing("strata::kernels::hadamard_rows"); +} +void hadamard_rows_f16(uint16_t*, int64_t, int64_t, int, const float*, void*) { + strata::metal_engine::missing("strata::kernels::hadamard_rows_f16"); +} +const float* hadamard_device_signs(const float*, int) { + strata::metal_engine::missing("strata::kernels::hadamard_device_signs"); + return nullptr; +} +} // namespace strata::kernels diff --git a/edge0-strata/metal/src/engine/window.cpp b/edge0-strata/metal/src/engine/window.cpp new file mode 100644 index 00000000..3ba6cdf2 --- /dev/null +++ b/edge0-strata/metal/src/engine/window.cpp @@ -0,0 +1,3334 @@ +// metal/src/engine/window.cpp - the decode window on Metal: Verifier::init/run/commit's Metal side (the one backend +// hook, src/core/verify.cpp under STRATA_GPU_METAL). Everything above it - the decode loop, MTP control, the serve +// protocol, the PLE table reads, the file tier - is the shared host code. +// +// The window is mac-kern's chain (metal/kernels, tests/token_chain.cpp) on the model GGUF itself: shard 1 is mmapped +// and wrapped no-copy as one MTLBuffer, every weight is read where it lies (unified memory: no upload, no expert +// cache), and one window is encoded as a dependency graph in one command buffer. Its own state lives here (the GDN +// recurrence and conv history, the int8 K/V and the indexer per QSA layer, the PLE history): the engine's +// SessionState buffers are not used by it, except the PLE token window (ss.ple_prev) and table. +// +// First version (MAC-PLAN Phase 3 start): windows of one token (no MTP drafts yet: --mtp-max-t 1), greedy. The PLE +// block's projections run on the GPU; its norms, gate and conv (10240 values) on the CPU between two command buffers +// (layer 0 | layers 1..47 + head). +#include "gguf_lite.h" +#include "mtl_rt.h" +#include "resource_path.h" +#include "ref_quant.h" +#include "arena_budget.h" +#include "stream.h" +#include "../../tests/coop_probe.h" + +#include "strata/core/layer.hpp" +#include "strata/core/conversation_cache.hpp" +#include "strata/core/latent_dense.hpp" +#include "strata/core/session.hpp" +#include "strata/kernels/ngram.hpp" +#include "strata/kernels/cpu/hadamard.hpp" +#include "strata/kernels/cpu/expert_layout.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace strata::gpu { uint64_t purge_device_allocations(size_t min_bytes); } // metal/src/gpu/runtime.cpp + +namespace strata::metal_engine { + +// The prompt GDN scan kernel, chosen by chip and prompt length (gdn_prompt.metal, gdn_scan_r; every choice is +// bit-identical to gdn_scanr_4: the token order and the arithmetic are the same, only the staged block size TB changes, which +// changes the threadgroup memory a core can hold at once). Measured by metal/tests/gdn_chunk_bench: +// Apple M4 (the base chip): gdn_scanr_4_8, 1.16x on the scan, 1.11x on the GDN prompt stage, T 512..32K +// Apple M5 Max: gdn_scanr_4_4 from T >= 4096 (1.04x at 4K, 1.10x at 16K); below that the default is as fast +// anything else (other chips, unknown names): gdn_scanr_4 as before +// STRATA_METAL_GDN_SCAN= forces a kernel (A/B and diagnosis). +static const char* gdn_scan_name(const std::string& device, int T) { + static const std::string forced = [] { + const char* e = std::getenv("STRATA_METAL_GDN_SCAN"); + const std::string s = e ? e : ""; + return (s == "gdn_scanr_4" || s == "gdn_scanr_4_8" || s == "gdn_scanr_4_4" || s == "gdn_scanr_4_2") ? s : std::string(); + }(); + if (!forced.empty()) return forced.c_str(); + if (device == "Apple M4") return "gdn_scanr_4_8"; + if (device == "Apple M5 Max" && T >= 4096) return "gdn_scanr_4_4"; + return "gdn_scanr_4"; +} + +// The M5 family's dense GEMM tiles for the prompt (bit-identical to the 64 x 64 tile: the tile only splits the output; measured on +// an M5 Max, T = 2048 / 4096): the GDN out projection (N 2560, K 6144) 64 x 128 (1.17 -> 0.99 ms), the wide input projections +// (N >= 6144) 128 x 64. Other chips keep 64 x 64 (M4: no gain). STRATA_METAL_GEMM_TILES=0 keeps the old tiles everywhere. +static void gemm_tile(const std::string& device, int Nn, int K, int& tm, int& tn) { + static const bool off = [] { const char* e = std::getenv("STRATA_METAL_GEMM_TILES"); return e && std::string(e) == "0"; }(); + if (off || device.rfind("Apple M5", 0) != 0) return; + if (Nn == 2560 && K == 6144) { tm = 64; tn = 128; } + else if (Nn >= 6144) { tm = 128; tn = 64; } +} +// The four GDN input projections (qkv, z, alpha, beta) as ONE GEMM with a 16480 float output row (M5 family: 3.09 -> 2.87 ms a layer at +// T = 2048; gdn_pre / gdn_post read the pieces at that stride, function constant 7). Bit-identical to the separate GEMMs. +// STRATA_METAL_GDN_FUSE=0 keeps the separate ones. +static bool gdn_fused_inputs(const std::string& device) { + static const bool off = [] { const char* e = std::getenv("STRATA_METAL_GDN_FUSE"); return e && std::string(e) == "0"; }(); + return !off && device.rfind("Apple M5", 0) == 0; +} + +namespace { + +int g_slot_cap = 0; // --batch N: the window keeps N + 1 sequences' running states side by side (set_batch_slots) + +constexpr int N = 2560, HC = 4, LR = 320, NFF = 640, TOPK = 10, NEXP = 512, MAXT = 8, NVOCAB_MAX = 248320; +constexpr int P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4, + P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE, P_IGU = 4, P_IDN = 8; +struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; +struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; +struct MoeArgs { + uint64_t gate_off, up_off, down_off, gu_stride, d_stride; + uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad; +}; +struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; }; +struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; }; +struct GlueArgs { uint32_t n, n_tok, n_embd, topk; }; +struct LatentRoutedArgs { + uint32_t n_tok, groups, hidden, latent, ff, topk, plan_token_entry; + uint32_t u_row, gu_row, down_row, up_off, down_off, expert_bytes; +}; +struct ArgArgs { uint32_t n_vocab, parts; }; +// the routed gate/up's i-quant grids from constant memory (QUANT_CTAB, function constant 23 of quant.metal): bitwise +// the threadgroup copy's values, 3-7% faster on the M5 Max; STRATA_METAL_CTAB=0 turns it off +bool metal_ctab() { + static const bool on = [] { + const char* e = std::getenv("STRATA_METAL_CTAB"); + return !(e && std::string(e) == "0"); + }(); + return on; +} + +int xkind(int t) { + switch (t) { + case ref::Q2_0: return 1; + case ref::Q4_K: case ref::Q5_K: return 2; + case ref::Q6_K: case ref::Q3_K: return 3; + default: return 0; + } +} + +// a prepared activation: 4 planes of MAXT rows of n, then the per-16 sums +struct Act { + mtl::Buffer b; + int n = 0; + size_t plane_bytes() const { return (size_t) 4 * MAXT * n; } +}; + +double now_ms() { return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); } + +} // namespace + +// the prompt chain's state (prompt_chain.inc): kernel libraries, the variant, buffers for chunks of C rows +struct PromptBufs { + int lib = -1, lgdn = -1, lqsa = -1; + bool k4 = false; + bool v3e = false; + bool fused = false; // the GDN input projections as one GEMM (gdn_fused_inputs) // the M5's expert kernels (v3): not for the trellis types, which have only v2 + int C = 0, maxrows = 0; + uint32_t expert_tm = 16; + double gpu_ms = 0, host_ms = 0, mtp_ms = 0; + int gu3[4] = {}, dn3[4] = {}; + mtl::Buffer R, bo, inj[2], xn, xnh, V, loh, mixed, mixh, qkv, z, al, be, yg, ygh, gqn, gkn, gvc, ggb, gkq, go, qfull, + kcur, vcur, iqr, ikr, qcur, qcurh, qidx, scores, qids, part, XQ, xqh, logits, ids, wts, bcnt, boff, eoff, tiles, meta, + stok, sw, pos, XS, HS, YE, yr, gsh, ush, hh, ys, sgv, WD, WU, Wqkv, Wz, Wal, Wbe, Wout, Wq, Wk, Wv, Wiq, Wik, Wr, + Wsg, Wshg, Wshu, Wshd, yf, Wcat, pemb, pembh, pkey, pval, keyn, qn, gated, normW, Wpk, Wpv, mixr; + mtl::Buffer lat_w, lat_proj, lat_pooled, lat_gu16; +}; + +// The vision path (--vision, GENI requests), from the shared host code: the image rows of the prompt being read +// (host rows by absolute position, null = embed the token; set by prompt_run for its duration) and the M-RoPE table +// (int32 [cells][3] (t, h, w); mrope_table_set: the host view of the device table the server fills per request). +const float* const* g_embd_rows = nullptr; +const int32_t* g_mrope_host = nullptr; +const int32_t* g_mrope_dev = nullptr; + +class Window { +public: + bool init(const strata::core::ModelGeometry& g, strata::core::SessionState& ss, const std::string& shard1, int max_t, + int64_t max_cells, std::string& err); + bool run(int T, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, std::string& err); + bool commit(int n_keep, std::string& err); + bool finish(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, double t0); + bool copy_logits(int t, float* host) const; + int64_t vocab() const { return n_vocab_; } + bool ready() const { return cx_ != nullptr; } + bool streaming() const { return stream_; } + bool latent() const { return latent_on_; } + void report_profile() const; + // --serve's conversation checkpoints: the window's running state (the GDN recurrences and conv histories, the + // PLE history, each QSA layer's indexer tail and tail-block key); the KV and pooled keys are positional + bool checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err); + bool checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err); + void reset(); // a new sequence (session_zero): the running state zeroed + // ---- batch slots (--batch N): each slot is another sequence's running state (GDN recurrence and conv history, the PLE + // history, each QSA layer's K/V, indexer tail and pooled keys). The live members always hold ONE sequence's state (the + // main one between calls); a slot's lives in a stash and is swapped in for its window (activate): the buffer objects stay + // where they are (the cached dispatches point at them), only their contents are exchanged. + bool slots_init(const std::vector& slots, std::string& err); + int n_slots() const { return sets_.empty() ? 0 : (int) sets_.size() - 1; } + int set_of(const strata::core::SessionState* s) const { + auto it = set_of_ss_.find(s); + return it == set_of_ss_.end() ? 0 : it->second; + } + // one window of S one-token rows, row t = slot rows[t] at position pos[t]: each runs and commits in its own state + bool run_slots(const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err); + // slot b's state <-> the main one (the GPU is idle between windows): the running state and the first `upto` cells of K/V + bool slot_copy(bool into_slot, int b, int64_t upto, bool kv_only, std::string& err); + void activate(int set); + void set_max_t(int t) { max_t_ = std::min(t, MAXT); } + // ---- the MTP draft layer (mtp.hpp's math on mac-kern's kernels; the MTP GGUF from tools/mtp_pack.py) + bool mtp_load(const std::string& gguf, int max_t, std::string& err); + bool mtp_prefill(const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err); + bool mtp_draft(int T, const int32_t* tokens, int64_t p, int a, int max_steps, int32_t* drafts, float* probs, + float min_p, int* n_drafts, std::string& err); + bool mtp_loaded() const { return mtp_on_; } + // the prompt through mac-kern's prompt kernels (prompt_chain.inc) + bool prompt(const int64_t* tokens, int64_t n, int64_t pos0, + const std::function& on_chunk, std::string& err); + +private: + const gguf::Tensor& T(const std::string& name) const { + const gguf::Tensor* t = f_.find(name); + if (!t) { std::fprintf(stderr, "strata metal: the model has no tensor %s\n", name.c_str()); std::exit(2); } + return *t; + } + const mtl::Pipeline* P(int lib, const std::string& fn, const mtl::Constants& c, const std::string& key) { + auto it = pcache_.find(key); + if (it == pcache_.end()) { + it = pcache_.emplace(key, cx_->pipeline(lib, fn, c)).first; + pkey_of_[&it->second] = key; + } + return &it->second; + } + void push(std::vector& ds, mtl::Dispatch d) { + d.barrier = !nb_; + nb_ = false; + // timing only (WRONG answers): STRATA_METAL_SKIP=a,b,... runs the dispatches whose pipeline name contains one + // of the words as empty grids (the marginal cost of a kernel class by leaving it out) + static const std::vector skip = [] { + std::vector v; + if (const char* e = std::getenv("STRATA_METAL_SKIP")) { + std::string s = e, w; + for (char c : s + ",") { if (c == ',') { if (!w.empty()) v.push_back(w); w.clear(); } else w += c; } + } + return v; + }(); + if (!skip.empty()) { + const std::string& k = pkey_of_[d.p]; + for (const auto& w : skip) + if (k.find(w) != std::string::npos) { d.indirect = nullptr; d.grid[0] = 0; break; } + } + mtl::require_fit(d, "engine dispatch (graph build)"); // a kernel Metal would skip silently: stop at the build, not at the first wrong answer + ds.push_back(std::move(d)); + } + std::map pkey_of_; + void gemv(std::vector& ds, const gguf::Tensor& t, const Act& x, const mtl::Buffer& y, size_t y_stride, + int ntok, size_t y_off = 0, const mtl::Buffer* wb = nullptr); + void xprep(std::vector& ds, const mtl::Buffer& x, const Act& out, int n); + void gr_read(std::vector& ds, const std::string& pre, bool apply, const Act& out, bool final_, int ntok); + void layer(std::vector& ds, int l, int ntok, int64_t pos0); + void ple_cpu(int T); + // the MTP layer: the input branches + attention HC read + K/V append for `n` rows (srcR rows, host tokens) at + // cells cell0..; with `full` (n == 1) the rest of the layer too and the draft head into mlm_/tokout_. The + // embedding rows go to memb_ rows erow..; tokens == nullptr: memb_ row erow already holds the embedding (put + // there by mtp_embed on the GPU). `slot`: the dense attention's cell list (one per draft step of a round). + void mtp_encode(std::vector& ds, const mtl::Buffer& srcR, size_t src_off, int n, const int32_t* tokens, + int64_t cell0, bool full, int erow = 0, int slot = 0); + // a draft round's step j on the GPU: the pick and its probability (dout_/dprob_ [j]), and with `embed` the picked + // token's embedding into memb_ row 0 (the next step's input) + void mtp_pick(std::vector& ds, int j, bool embed); + void gr_read_m(std::vector& ds, const std::string& pre, mtl::Buffer& Rb, const mtl::Buffer& bo, + const mtl::Buffer& inj_in, const mtl::Buffer& inj_out, bool apply, const Act& out, bool final_, int ntok); + const gguf::Tensor& TM(const std::string& name) const { + const gguf::Tensor* x = fm_.find(name); + if (!x) { std::fprintf(stderr, "strata metal: the MTP model has no tensor %s\n", name.c_str()); std::exit(2); } + return *x; + } + size_t NM(const std::string& name) const { // the folded (1 + w) copy of an F32 norm of the MTP model + auto it = mnorm_off_.find(name); + if (it == mnorm_off_.end()) { std::fprintf(stderr, "strata metal: no MTP norm %s\n", name.c_str()); std::exit(2); } + return it->second; + } + + strata::core::SessionState* ss_ = nullptr; // the session whose sequence the live state is (the main one, or a batch slot's) + strata::core::SessionState* main_ss_ = nullptr; // the main session: it owns the PLE table (read by every slot) + gguf::File f_; + strata::core::LatentDense latent_dense_; + strata::core::LatentA latent_experts_; + std::unique_ptr cx_; + mtl::Buffer W_; + bool latent_on_ = false; + bool latent_grouped_u_ = false; + bool latent_fast_verify_ = false; + bool fast_verify(int ntok) const { + return latent_fast_verify_ && win_head_ && !rows_ && ntok > 1; + } + int latent_lib_ = -1; + mtl::Buffer latent_du_, latent_a_, latent_labels_, latent_inv_, latent_d_, latent_down_; + mtl::Buffer latent_group_active_; + LatentRoutedArgs latent_args_{}; + bool latent_init(std::string& err); + void latent_routed(std::vector& ds, int layer, int ntok); + // a model wrapped in chunks (mtl::Context::wrap): the routed experts are bound at the layer's first expert tensor, + // their offsets made relative to it; returns that binding offset (0 when W_ is one buffer or the experts stream) + size_t moe_rebase(uint64_t& g, uint64_t& u, uint64_t& d) const { + if (stream_ || !W_.chunks) return 0; + const uint64_t b = std::min({g, u, d}); + g -= b; u -= b; d -= b; + return (size_t) b; + } + int lib_ = -1, lgdn_ = -1, lqsa_ = -1, lsmp_ = -1, lple_ = -1; + std::map pcache_; + bool nb_ = false; + int max_t_ = 1; + int64_t max_cells_ = 0; + int64_t n_vocab_ = 0; + float eps_ = 1e-6f; + // buffers + mtl::Buffer R_, bo_, inj_[2], xn_, lo_, qkv_, z_, al_, be_, yg_, logits_, ids_, wts_, plan_, h_, ye_, gsh_, ush_, + ysh_, sg_, qfull_, kcur_, vcur_, iqr_, ikr_, qcur_, qidx_, scores_, qids_, part_, cs_, lm_, apv_, api_, tokout_, + emb_, pkey_, pval_, qkvL_, alL_, beL_, zL_, ydump_, keyn_, qn_, gated_, normW_, pleHist_; + Act XA_, XF_, XY_, XQ_, XH_, XE_; + int cur_inj_ = 0; + // state + std::vector gst_, ghist_[2], kq_, vq_, ksc_, vsc_, tail_, dead_, pooled_; + // --batch N: every one of those is a view of set 0 in a bigger buffer holding all N + 1 sets (the slots' running states side + // by side), so ONE dispatch can address any slot (rinfo_: the strides and each row's slot / cell / history parity) + bool bigs_ok_ = false; + std::vector gstBig_, hbBig_, kqBig_, vqBig_, kscBig_, vscBig_, tailBig_, deadBig_, pooledBig_; + mtl::Buffer pleBig_, rinfo_; + size_t st_gst_ = 0, st_hist_ = 0, st_kq_ = 0, st_sc_ = 0, st_tail_ = 0, st_dead_ = 0, st_pooled_ = 0, st_ple_ = 0; // bytes between sets + bool rows_ = false; // the window being built / run holds one token of each of rows_n_ DIFFERENT sequences (batch slots) + int rows_n_ = 0; + strata::core::SessionState* rows_ss_[MAXT] = {}; + int rows_set_[MAXT] = {}; + static mtl::Buffer view_of(const mtl::Buffer& big, size_t off, size_t bytes) { + mtl::Buffer v = big; + v.ptr = (char*) big.ptr + off; + v.bias = big.bias + off; + v.bytes = bytes; + return v; + } + bool run_rows(const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err); + bool finish_rows(int S, const int32_t* tokens, int32_t* out, double t0); + struct SeqSet { + std::vector gst, gh0, gh1, kq, vq, ksc, vsc, tail, dead, pooled; + mtl::Buffer pleHist; + std::vector ple_hist; + std::vector last_norm; + int hp = 0, pend_fold = 0, last_T = 0; + bool last_head = true; + int32_t last_tokens[MAXT] = {}; + strata::core::SessionState* ss = nullptr; + }; + std::vector> sets_; // [0] the spare the main state is exchanged with, [b + 1] slot b + std::map set_of_ss_; + int active_ = 0; + void swap_state(SeqSet& o); + void alloc_set(SeqSet& o, bool real, int idx = 0); + int hp_ = 0; // which of the two history buffers holds the committed conv history (all layers) + // commit(n) of a window of several tokens: the GDN fold of the kept prefix waits for the next window and runs at + // the start of its command buffer (one submit and wait less a round). Nothing else reads the GDN state between + // the two (the drafter reads the window's rows and its own caches). STRATA_METAL_LAZY_FOLD=0: at commit. + int pend_fold_ = 0; + // gdn_step's partial commit (mac-kern 56526aed, commit = 0x100 | k): a verify window of T > 1 commits its first + // row in place (it is always kept), so commit(1) needs no fold and commit(n) folds rows 1..n-1 only; a window + // without the head (the prompt) commits all its rows in the step. Opt-in (STRATA_METAL_PARTIAL=1): at --spec 2 + // it measured 0.3-0.8% slower on the M5 Max (the in-place state write costs every window, the fold it saves + // only the rounds with no accepted draft). Off: the window commits nothing and commit(n) folds n rows. + bool partial_ = std::getenv("STRATA_METAL_PARTIAL") && std::string(std::getenv("STRATA_METAL_PARTIAL")) == "1"; + bool win_head_ = true; // the window being encoded has the head + bool last_head_ = true; // the last window had it + // token_chain's --fuse (mac-kern): the GDN step writes the prepared y planes itself (no xprep), and a window whose + // context is within the selection budget skips the QSA score and top-k (the cell list is the identity: qids_ + // holds it from init, restored after a window that selected). STRATA_METAL_FUSE=0: off. + bool fuse_ = !(std::getenv("STRATA_METAL_FUSE") && std::string(std::getenv("STRATA_METAL_FUSE")) == "0"); + bool qdirty_ = false; + // mac-kern's router_route (3f9ca996): the BF16 router GEMV and moe_route in one dispatch, beside the shared + // expert's GEMVs (bitwise). Opt-in (STRATA_METAL_ROUTE=1): in this window, where the router GEMV already runs + // beside the shared GEMVs, it measured no faster (T2 GPU 13.035 vs 13.027 ms; the fused kernel alone 0.615 ms vs + // 0.566 for the two) and 0.3-0.5% slower end to end on the M5 Max. + bool routef_ = std::getenv("STRATA_METAL_ROUTE") && std::string(std::getenv("STRATA_METAL_ROUTE")) == "1"; + mtl::Buffer rcounter_; + // STRATA_METAL_KPROF=T (test tooling): once, after the 20th cached window of T tokens with the head, the GPU time + // of each kernel class alone (the window's own dispatches of that pipeline + weight role, run by themselves) + void kprof(const std::vector& ds, int Tn); + void fold(std::vector& ds, int n_keep, int hp); + // the encoded window per (T, head, hp_): built once, then only the positions are patched (the host's encode cost) + struct Cached { + std::vector ds; + std::vector cuts; // SSD streaming: the segment ends (one per layer, before its experts) + std::vector> qsa; // (dispatch, kind): 0 QsaArgs n_tok = T, 1 the merge (MAXT), 2 the scores + }; + std::map wcache_; + std::vector>* rec_ = nullptr; + // ---- SSD streaming (stream.h): the dense weights in W_, the routed experts in the arena + std::unique_ptr st_; + bool stream_ = false, st_pred_ = false; + std::vector* cuts_ = nullptr; + mtl::Buffer predL_, predL2_; // 48 x MAXT x 512: layer l + 1's (l + 2's) router logits from layer l's FFN input + int seg_T_ = 0; + bool stream_arena(std::string& err); + double run_ds(Cached& c); + // prompts on the matrix units (prompt_chunk below) + struct PromptGemm; + PromptGemm* pg_ = nullptr; + const int32_t* cut_ids_ = nullptr; // the routed ids the handshake reads (a prompt chunk's; else ids_) + const mtl::Buffer* mtp_src_ = nullptr; // the rows the draft layer's prompt K/V read (a prompt chunk's; else R_) + double pg_ms_ = 0; + int64_t pg_tok_ = 0; +public: + bool prompt_gemm_on(int64_t n) const; + bool prompt_gemm_init(std::string& err); + bool prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::string& err); + int prompt_chunk_rows() const; + void prompt_region(bool hold) { if (stream_ && st_ && st_->ready()) st_->hold_region(hold); } + void stream_report(FILE* f) const { + if (!st_) return; + st_->report(f); + for (int t = 1; t <= MAXT; ++t) + if (win_n_[t]) + std::fprintf(f, "stream windows T=%d %d (%s): %.2f ms each (GPU incl. the host's turns %.2f ms)\n", t, win_n_[t], + t == 8 ? "prompt rows" : "decode", win_ms_[t] / win_n_[t], win_gpu_[t] / win_n_[t]); + } + int win_n_[MAXT + 1] = {}; + double win_ms_[MAXT + 1] = {}, win_gpu_[MAXT + 1] = {}; + void at_cut(int s) { + st_->at_layer(s, seg_T_, TOPK, cut_ids_ ? cut_ids_ : ids_.as(), st_pred_ && !cut_ids_ ? predL_.as() + (size_t) s * MAXT * 512 : nullptr, + predL2_.ptr && !cut_ids_ ? predL2_.as() + (size_t) s * MAXT * 512 : nullptr); + } +private: + std::vector ple_hist_; // NG_HIST x hc_dim normalized rows, row-fastest + bool ple_on_ = false; + int last_T_ = 0; + double last_gpu_ms_ = 0; + int last_n_disp_ = 0; + int32_t last_tokens_[MAXT] = {}; + std::vector last_norm_; // the window's normalized PLE rows (T x hc_dim), for commit + // STRATA_METAL_LOGPOS= (test tooling, one-token windows): per decode position the log-probability of the + // token that came next (the next window's input: with --spec-follow, the forced text) and the top 64, in the + // KL harness's row format (phaseA-tests cuda_p0/cuda_kl.py load_lp): pos, target, lp, argmax, 4 x 0, id:lp ... + std::FILE* logpos_ = nullptr; + bool lp_pend_ = false; + int64_t lp_pos_ = 0, lp_first_ = -1; + int32_t lp_arg_ = 0; + double lp_lse_ = 0.0; + std::vector lp_row_; + std::vector> lp_top_; + void logpos(int Tn, const int32_t* tokens, int64_t pos0, bool head); + // STRATA_METAL_PROFILE=1: where a window's time goes, per window size (printed when the window goes away) + struct Prof { int64_t n = 0; double pre = 0, rg = 0, gpu = 0, fin = 0, all = 0; }; + std::map prof_; + Prof cprof_, dprof_; // commit (folds), the draft rounds + double tp_pre_ = 0, tp_rg_ = 0; +public: + ~Window(); + void prof_add(bool draft, double ms) { Prof& x = draft ? dprof_ : cprof_; ++x.n; x.all += ms; } +private: + // the Q8_0 copies of the hyper-connection matrices (the default; STRATA_HC_Q8=0: the BF16 originals) + // Gyro (agentionai's Hadamard-folded trellis experts, native_experts.txt v5 + hadamard.txt): the gate/up input is + // H (s * x) and the down input H h, per block of 128; XR_ holds the rotated rows, hsx_/hsh_ the +/-1 signs + bool had_on_ = false; + int had_block_ = 128; + bool gyro_dense_ = false; // a trellis-expert file: its dense K-quants get their own GEMV tiles (gemv() below) + bool had_gu(int l) const { return had_on_ && strata::kernels::cpu::hadamard_spec().gu(l); } + bool had_dn(int l) const { return had_on_ && strata::kernels::cpu::hadamard_spec().down(l); } + mtl::Buffer XR_, hsx_, hsh_; + // one rotation dispatch: rows of `width` values (a multiple of 128) in place or from x to y; f16 = half rows + void had_rows(std::vector& ds, int lib, const std::string& pfx, const mtl::Buffer& x, size_t x_off, + const mtl::Buffer& y, size_t y_off, int rows, int width, const mtl::Buffer* signs, bool f16); + // the prompt's h rows (half, padded rows of the expert tiles, their count in meta[8]): rotated in place, one + // threadgroup a row through the indirect grid + void had_rows_ind(std::vector& ds, int lib, const std::string& pfx, const mtl::Buffer& h, const mtl::Buffer& meta); + bool hc_q8_ = false; + mtl::Buffer Q8_; + std::map q8off_; + // MTP + bool mtp_on_ = false; + int lmtp_ = -1; + gguf::File fm_; + mtl::Buffer Wm_, mnorm_, mR_, memb_, men_, me2_, mhn_, mh2_, midx_, mbo_, minj_[2], mqids_, mlm_; + std::map mnorm_off_; + std::vector mkq_, mvq_, mksc_, mvsc_, mtail_, mdead_, mpooled_; + // the draft head over a token subset (data/draft_vocab.bin, as the CUDA drafter): its rows gathered from the head + mtl::Buffer mhead_; + gguf::Tensor mhead_t_; + std::vector dvocab_; + // the draft layer's BF16 matrices in Q8_0 (half the bytes a draft step reads; drafts only: the verify window + // decides every token). STRATA_METAL_MTP_Q8=0 keeps BF16. + mtl::Buffer mq8_; + std::map mq8t_; + const mtl::Buffer& MB(const std::string& name) const { return mq8t_.count(name) ? mq8_ : Wm_; } + const gguf::Tensor& TQ(const std::string& name) const { + auto it = mq8t_.find(name); + return it != mq8t_.end() ? it->second : TM(name); + } + Act XEn_, XHn_; + PromptBufs pb_; + mtl::Buffer mpemb_; // the drafter's prompt embedding rows (mtp_prefill, one command buffer) + size_t mpemb_rows_ = 0; + const mtl::Buffer* membp_ = nullptr; // mtp_encode's embedding rows (nullptr: memb_) + const mtl::Buffer* mpsrc_ = nullptr; // the drafter's prompt rows (nullptr: R_) + bool coop_k4(); + // the vision path: the embedding row of an absolute position (an image row, or null), the M-RoPE table rows a + // window reads copied into mrope_ (qsa_prep's buffer 20, function constant 24) + const float* img_row(int64_t pos) const { return g_embd_rows ? g_embd_rows[pos] : nullptr; } + mtl::Buffer mrope_; + // the PLE conv kernel as F16 (ple_conv_add reads half): the file's own F16 tensor, or an F32 one (Unsloth's) + // narrowed with round-to-nearest-even as the host loader does + mtl::Buffer plekw_; + bool plekw_own_ = false; + mtl::Arg kw_arg() { return plekw_own_ ? mtl::Arg::B(plekw_) : mtl::Arg::B(W_, T("blk.1.ple_conv1d.weight").offset); } + const uint16_t* kw_host() { + return plekw_own_ ? plekw_.as() + : (const uint16_t*) ((const uint8_t*) W_.ptr + T("blk.1.ple_conv1d.weight").offset); + } + bool mrope_on() const { return g_mrope_host != nullptr; } + void mrope_sync(int64_t c0, int64_t n) { + if (!mrope_on()) return; + const int64_t a = std::max(0, c0 & ~(int64_t) 3), b = std::min(max_cells_, c0 + n); + if (b > a) std::memcpy(mrope_.as() + a * 3, g_mrope_host + a * 3, (size_t) (b - a) * 12); + } + mtl::Constants qsa_c() const { return mrope_on() ? mtl::Constants().b(24, true) : mtl::Constants(); } + std::string qsa_k() const { return mrope_on() ? "qsa_prep#mrope" : "qsa_prep"; } + int m5_layout_ = -1; // the M5's matmul2d layout (coop_k4), asked once: the M4's tiles differ + bool m5_layout() { if (m5_layout_ < 0) m5_layout_ = coop_k4() ? 1 : 0; return m5_layout_ == 1; } + bool pc_setup(int C); + bool pc_chunk(const int64_t* tokens, int nT, int64_t pos0, std::string& err); + // the draft round in one command buffer (msl/mtp_chain.metal): STRATA_METAL_MTP_CHAIN=0 runs a command buffer + // per step with the host's pick and embedding (as before) + int lchain_ = -1; + bool chain_ = false; + mtl::Buffer dout_, dprob_, posv_, dmap_; +}; + +namespace { +// A repack (metal_repack.py: .repack.lock) or setup's model preparation (.preparing.lock) beside the GPU's file holds a lock +// with its pid. The repack and the pack move 100+ GB through the disk and one holds up to 20 GB of memory: an engine of 70-90 GB +// beside them is what froze the 128 GB M5 Max twice (watchdog panics, 2026-10-13). Returns the pid of a live one, else 0. +// STRATA_METAL_IGNORE_REPACK=1 starts anyway (your choice). +long repack_pid_running(const std::string& model_file) { + namespace fs = std::filesystem; + std::error_code ec; + const fs::path dir = fs::path(model_file).parent_path().empty() ? fs::path(".") : fs::path(model_file).parent_path(); + for (const auto& e : fs::directory_iterator(dir, ec)) { + const std::string n = e.path().filename().string(); + const bool lock = n == ".preparing.lock" || (n.size() > 12 && n.compare(n.size() - 12, 12, ".repack.lock") == 0); + if (!lock) continue; + std::FILE* f = std::fopen(e.path().string().c_str(), "r"); + long pid = 0; + if (f) { if (std::fscanf(f, "%ld", &pid) != 1) pid = 0; std::fclose(f); } + if (pid > 0 && pid != (long) getpid() && (kill((pid_t) pid, 0) == 0 || errno == EPERM)) return pid; + } + return 0; +} +} // namespace + +bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionState& ss, const std::string& shard1, + int max_t, int64_t max_cells, std::string& err) { + ss_ = &ss; + main_ss_ = &ss; + max_t_ = std::min(max_t, MAXT); + max_cells_ = max_cells; + if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert != NEXP || g.n_layers != 48) { + err = "metal window: the geometry differs from Qwen3.8-Flash-Next's"; + return false; + } + eps_ = 1e-6f; + const double t0 = now_ms(); + if (const long rp = repack_pid_running(shard1); rp > 0 && !(std::getenv("STRATA_METAL_IGNORE_REPACK") && std::string(std::getenv("STRATA_METAL_IGNORE_REPACK")) == "1")) { + err = "metal window: a model repack or setup's preparation is running here (pid " + std::to_string(rp) + "): the engine is not started " + "beside it (it froze a 128 GB Mac twice); wait for it to finish, or set STRATA_METAL_IGNORE_REPACK=1"; + return false; + } + if (!f_.open(shard1, err)) return false; + cx_ = std::make_unique(); + if (!latent_init(err)) return false; + stream_ = !latent_on_ && ExpertStream::wanted(f_, *cx_); + if (stream_) { + // the model does not fit: the dense weights are read into memory below, the experts stream from the SSD + st_ = std::make_unique(); + if (!st_->open(f_, err)) return false; + } else { + { // read ahead the weights, not the n-gram table: in a Gyro file it is 28.8 GB in the middle of the one file, read by the host from + // the SSD a few rows at a time (src/ngram/ple_reader.cpp), and must not fill the page cache + const gguf::Tensor* ple = f_.find("per_layer_token_embd.weight"); + if (ple == nullptr) { + madvise(f_.map, f_.map_bytes, MADV_WILLNEED); + } else { + const size_t page = 16384; + const size_t lo = (size_t) ple->offset / page * page, hi = ((size_t) (ple->offset + ple->bytes) + page - 1) / page * page; + if (lo > 0) madvise(f_.map, lo, MADV_WILLNEED); + if (hi < f_.map_bytes) madvise((uint8_t*) f_.map + hi, f_.map_bytes - hi, MADV_WILLNEED); + } + } + // the chunks of a model bigger than one Metal buffer overlap by the largest range bound at one offset: a tensor, + // or a layer's routed experts (gate, up and down, bound together at the first of them: moe_rebase) + size_t span = 0; + // (not the n-gram table: the host reads it from the SSD, the GPU never binds it - in a Gyro file it is 28.8 GB of the one file) + for (const auto& kv : f_.tensors) + if (kv.first != "per_layer_token_embd.weight") span = std::max(span, kv.second.bytes); + for (int l = 0; l < 48; ++l) { + const std::string b = "blk." + std::to_string(l) + "."; + const gguf::Tensor *g = f_.find(b + "ffn_gate_exps.weight"), *u = f_.find(b + "ffn_up_exps.weight"), + *d = f_.find(b + "ffn_down_exps.weight"); + if (!g || !u || !d) continue; + const uint64_t lo = std::min({g->offset, u->offset, d->offset}); + const uint64_t hi = std::max({g->offset + g->bytes, u->offset + u->bytes, d->offset + d->bytes}); + span = std::max(span, hi - lo); + } + W_ = cx_->wrap(f_.map, f_.map_bytes, span); + if (W_.handle == nullptr) { err = "metal window: cannot wrap " + shard1 + " as a Metal buffer"; return false; } + } + n_vocab_ = (int64_t) T("output.weight").shape[1]; + if (const gguf::Tensor* e0 = f_.find("blk.0.ffn_gate_exps.weight")) gyro_dense_ = e0->type >= 144; + lib_ = cx_->compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "router.metal", "gr.metal", "gr2.metal", + "glue.metal"}); + rcounter_ = cx_->alloc(16); + lgdn_ = cx_->compile({"gdn.metal"}, "#define NO_GGML_COMMON 1\n"); + lqsa_ = cx_->compile({"qsa.metal", "qsa2.metal"}, "#define NO_GGML_COMMON 1\n"); + lsmp_ = cx_->compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n"); + lple_ = cx_->compile({strata_metal::resource("STRATA_METAL_ENGINE_KERNELS", "engine-msl", STRATA_METAL_ENGINE_KERNELS_DEFAULT) + "/ple.metal"}, "#define NO_GGML_COMMON 1\n"); + + const int MAXC = (int) std::max(4096, (max_cells + 64) / 4 * 4 + 64), MB = MAXC / 4 + 2; + max_cells_ = MAXC; + auto A = [&](size_t bytes) { return cx_->alloc(bytes); }; + R_ = A(4 * MAXT * HC * N); bo_ = A(4 * MAXT * N); inj_[0] = A(4 * MAXT * HC); inj_[1] = A(4 * MAXT * HC); + xn_ = A(4 * MAXT * HC * N); lo_ = A(4 * MAXT * LR); qkv_ = A(4 * MAXT * 10240); z_ = A(4 * MAXT * 6144); + al_ = A(4 * MAXT * 48); be_ = A(4 * MAXT * 48); yg_ = A(4 * MAXT * 6144); logits_ = A(4 * MAXT * 512); + ids_ = A(4 * MAXT * TOPK); wts_ = A(4 * MAXT * TOPK); plan_ = A(4 * P_INTS); h_ = A(4 * MAXT * TOPK * NFF); + ye_ = A(4 * MAXT * TOPK * N); gsh_ = A(4 * MAXT * NFF); ush_ = A(4 * MAXT * NFF); ysh_ = A(4 * MAXT * N); + sg_ = A(4 * MAXT * 4); qfull_ = A(4 * MAXT * 12288); kcur_ = A(4 * MAXT * 512); vcur_ = A(4 * MAXT * 512); + iqr_ = A(4 * MAXT * 512); ikr_ = A(4 * MAXT * 128); qcur_ = A(4 * MAXT * 24 * 256); qidx_ = A(4 * MAXT * 4 * 128); + scores_ = A(4 * (size_t) MAXT * MB); qids_ = A(4 * MAXT * 2051); part_ = A(4 * (size_t) MAXT * 65 * 24 * 258); + cs_ = A(8 * (size_t) (MAXC + 16) * 32); lm_ = A(4 * (size_t) MAXT * n_vocab_); apv_ = A(4 * MAXT * 120); + api_ = A(4 * MAXT * 120); tokout_ = A(4 * MAXT); emb_ = A(4 * MAXT * N); pkey_ = A(4 * MAXT * 10240); + pval_ = A(4 * MAXT * N); + // the GDN layers' window inputs, kept per layer so commit() can fold the accepted prefix into the state + qkvL_ = A(4 * (size_t) 36 * MAXT * 10240); alL_ = A(4 * (size_t) 36 * MAXT * 48); beL_ = A(4 * (size_t) 36 * MAXT * 48); + zL_ = A(4 * (size_t) 36 * MAXT * 6144); ydump_ = A(4 * MAXT * 6144); + keyn_ = A(4 * MAXT * 10240); qn_ = A(4 * MAXT * 10240); gated_ = A(4 * MAXT * 10240); normW_ = A(4 * MAXT * 10240); + mrope_ = A(12 * (size_t) (MAXC + 16)); + bigs_ok_ = g_slot_cap > 0; + const size_t NSET = 1 + (size_t) std::max(0, g_slot_cap); + auto R256 = [](size_t b) { return (b + 255) / 256 * 256; }; + st_gst_ = R256(4 * 128 * 48 * 128); st_hist_ = R256(4 * 10240 * 3); st_kq_ = R256((size_t) MAXC * 2 * 256); + st_sc_ = R256(2 * (size_t) MAXC * 2 * 4); st_tail_ = R256(4 * 4 * 128); st_dead_ = R256(4 * 128); + st_pooled_ = R256(4 * (size_t) MB * 128); st_ple_ = R256(4 * 9 * 10240); + // the set-0 view of a new buffer of `bytes` per set (or a plain buffer when there are no batch slots) + auto SV = [&](std::vector& bigs, size_t bytes, size_t stride, size_t per = 1, size_t which = 0) { + if (!bigs_ok_) { bigs.push_back(mtl::Buffer()); return A(bytes); } + if (which == 0) bigs.push_back(A(stride * per * NSET)); + return view_of(bigs.back(), which * stride, bytes); + }; + if (bigs_ok_) { pleBig_ = A(st_ple_ * NSET); pleHist_ = view_of(pleBig_, 0, 4 * 9 * 10240); } else pleHist_ = A(4 * 9 * 10240); + rinfo_ = A(4 * (8 + 4 * MAXT)); + auto act = [&](int n) { Act a; a.n = n; a.b = A(4 * ((size_t) 4 * MAXT * n + (size_t) MAXT * n / 16)); return a; }; + XA_ = act(N); XF_ = act(N); XY_ = act(6144); XQ_ = act(6144); XH_ = act(NFF); XE_ = act(N); + { // Gyro (native_experts.txt v5 + hadamard.txt): the rotation's signs and the rotated gate/up input + const auto& hs = strata::kernels::cpu::hadamard_spec(); + had_on_ = hs.any(); + if (had_on_) { + had_block_ = hs.block; + if (had_block_ != 128) { + err = "metal window: the pack's Hadamard block is " + std::to_string(had_block_) + "; the Metal rotation takes 128"; + return false; + } + XR_ = A(4 * MAXT * N); + auto upload = [&](int w, mtl::Buffer& b) { + if (const float* sg = strata::kernels::cpu::hadamard_signs(w)) { + b = A(4 * (size_t) w); + std::memcpy(b.ptr, sg, 4 * (size_t) w); + } + }; + upload(N, hsx_); + upload(NFF, hsh_); + std::fprintf(stderr, "strata metal: Hadamard-folded experts (block %d): the gate/up and down inputs are rotated on the GPU\n", + had_block_); + } + } + // RoPE: cos/sin [pos][32] for the first 64 dims, NEOX pairs (i, i + 32), in double as build_rope_table + const double base = f_.meta_num.count("qwen4exp.rope.freq_base") ? f_.meta_num["qwen4exp.rope.freq_base"] : 1e7; + for (int p = 0; p < MAXC + 16; ++p) + for (int i = 0; i < 32; ++i) { + const double th = (double) p * std::pow(base, -2.0 * i / 64.0); + cs_.as()[(p * 32 + i) * 2] = (float) std::cos(th); + cs_.as()[(p * 32 + i) * 2 + 1] = (float) std::sin(th); + } + for (int t = 0; t < MAXT; ++t) + for (int c = 0; c < 2051; ++c) qids_.as()[t * 2051 + c] = c; + for (int l = 0; l < 48; ++l) { + if (l % 4 != 3) { + gst_.push_back(SV(gstBig_, 4 * 128 * 48 * 128, st_gst_)); + ghist_[0].push_back(SV(hbBig_, 4 * 10240 * 3, st_hist_, 2, 0)); // [set][parity]: the two parities side by side + ghist_[1].push_back(SV(hbBig_, 4 * 10240 * 3, st_hist_, 2, 1)); + } else { + kq_.push_back(SV(kqBig_, (size_t) MAXC * 2 * 256, st_kq_)); + vq_.push_back(SV(vqBig_, (size_t) MAXC * 2 * 256, st_kq_)); + ksc_.push_back(SV(kscBig_, 2 * (size_t) MAXC * 2 * 4, st_sc_)); + vsc_.push_back(SV(vscBig_, 2 * (size_t) MAXC * 2 * 4, st_sc_)); + tail_.push_back(SV(tailBig_, 4 * 4 * 128, st_tail_)); + dead_.push_back(SV(deadBig_, 4 * 128, st_dead_)); + pooled_.push_back(SV(pooledBig_, 4 * (size_t) MB * 128, st_pooled_)); + } + } + // ---- the HC read in Q8_0 (the owner's default on Metal, as STRATA_HC_Q8 on gfx12): the BF16 down / up / inject + // matrices quantized once at load (ggml's q8_0 rule, ref_quant), read by gr_down_q8 / gr_up_q8 + { + const char* e = std::getenv("STRATA_HC_Q8"); + hc_q8_ = !(e && std::string(e) == "0"); + } + { + const double tq = now_ms(); + size_t tot = 0; + std::vector names; + for (const auto& kv : f_.tensors) { + const gguf::Tensor& x = kv.second; + // BF16 (when the Q8_0 read is on) and, whatever the setting, the Gyro file's Q6_K / IQ4_NL ones: the decode window's + // hyper-connection kernels read BF16 or Q8_0 only + const bool conv = (x.type == ref::BF16 && hc_q8_) || x.type == ref::Q6_K || x.type == ref::IQ4_NL; + if (!conv || x.name.find("hc_") == std::string::npos || x.name.find("_norm") != std::string::npos) continue; + uint64_t ne = 1; + for (uint64_t d : x.shape) ne *= d; + q8off_[x.name] = tot; + names.push_back(x.name); + tot += ne / 32 * 34; + } + Q8_ = cx_->alloc((tot + 4095) / 4096 * 4096 + 4096); // + a page: a file with Q8_0 HC already has no copies + std::vector th; + const unsigned nth = std::max(1u, std::min(16u, std::thread::hardware_concurrency())); + for (unsigned k = 0; k < nth; ++k) + th.emplace_back([&, k] { + std::vector tmp; + std::vector raw; + for (size_t i = k; i < names.size(); i += nth) { + const gguf::Tensor& x = *f_.find(names[i]); + uint64_t ne = 1; + for (uint64_t d : x.shape) ne *= d; + tmp.resize(ne); + const uint8_t* src = (const uint8_t*) f_.map + x.offset; + if (stream_) { // F_NOCACHE reads: the page cache is kept for the arena + raw.resize(x.bytes); + if (!st_->read(raw.data(), x.bytes, x.offset)) { std::fprintf(stderr, "strata metal: read failed\n"); std::exit(6); } + src = raw.data(); + } + ref::dequant(x.type, src, tmp.data(), ne); + // a Q6_K / IQ4_NL matrix (Gyro) goes through BF16 first: the form the pack gives the CUDA engine + // (tools/iq_pack.py --compat-bf16), so the Q8_0 copy below is made from the same numbers + if (x.type != ref::BF16) + for (uint64_t i = 0; i < ne; ++i) tmp[i] = ref::bf16_to_f32(ref::f32_to_bf16(tmp[i])); + ref::quantize_q8_0(tmp.data(), ne, Q8_.as() + q8off_[names[i]]); + } + }); + for (auto& x : th) x.join(); + std::fprintf(stderr, "strata metal: HC read in Q8_0 (STRATA_HC_Q8=0: BF16): %zu matrices, %.0f MB, made in %.0f ms\n", + names.size(), tot / 1e6, now_ms() - tq); + } + if (stream_) { + std::vector skip; // the BF16 originals of the Q8_0 HC copies are not read again + if (hc_q8_) + for (const auto& kv : q8off_) skip.push_back(kv.first); + if (!st_->load_dense(f_, *cx_, skip, W_, err)) return false; + const char* e = std::getenv("STRATA_METAL_PREFETCH"); + st_pred_ = !(e && std::atoi(e) <= 0); + if (st_pred_) predL_ = A(4 * (size_t) 48 * MAXT * 512); + if (st_pred_ && st_->ahead2() > 0) predL2_ = A(4 * (size_t) 48 * MAXT * 512); + } + ple_on_ = ss.ple.ready() && f_.find("blk.1.ple_key.weight") != nullptr; + if (ple_on_ && T("blk.1.ple_conv1d.weight").type == ref::F32) { + const gguf::Tensor& k = T("blk.1.ple_conv1d.weight"); + const size_t ne = k.bytes / 4; + plekw_ = cx_->alloc((ne * 2 + 4095) / 4096 * 4096); + const float* src = (const float*) ((const uint8_t*) W_.ptr + k.offset); + uint16_t* dst = plekw_.as(); + for (size_t i = 0; i < ne; ++i) dst[i] = ref::f32_to_f16(src[i]); + plekw_own_ = true; + } + ple_hist_.assign((size_t) strata::kernels::NG_HIST * strata::kernels::NG_HC_DIM, 0.0f); + last_norm_.assign((size_t) MAXT * strata::kernels::NG_HC_DIM, 0.0f); + std::fprintf(stderr, "strata metal: decode window on %s: %s %s (%.1f GB), %d cells, PLE %s, ready in %.0f ms\n", + cx_->device_name().c_str(), shard1.c_str(), stream_ ? "streamed from the SSD (experts)" : "wrapped", + f_.file_bytes / 1e9, MAXC, ple_on_ ? "on" : "OFF", now_ms() - t0); + return true; +} + +bool Window::latent_init(std::string& err) { + const auto& layout = strata::kernels::cpu::expert_layout(); + if (!layout.native || layout.fmt.empty() || layout.fmt.front().n_embd == N) return true; + constexpr uint32_t groups = 8, layers = 48; + const auto declared_width = layout.fmt.front().n_embd; + if (declared_width != 1280 && declared_width != 1792) { + err = "metal latent: expected latent width 1280 or 1792"; + return false; + } + const auto width = uint32_t(declared_width); + const char* enabled = std::getenv("STRATA_METAL_LATENT"); + const char* dense = std::getenv("STRATA_LATENT_DENSE"); + const char* experts = std::getenv("STRATA_LATENT_A"); + if (!enabled || std::strcmp(enabled, "1") || !dense || !*dense || !experts || !*experts) { + err = "metal latent: set STRATA_METAL_LATENT=1, STRATA_LATENT_DENSE and STRATA_LATENT_A"; + return false; + } + if (layout.fmt.size() != layers || layout.n_expert != NEXP) { + err = "metal latent: expected 48 layers and 512 experts"; + return false; + } + const size_t gu_row = ref::row_bytes(ref::Q3_K, width); + const size_t down_row = ref::row_bytes(ref::IQ4_NL, NFF); + const size_t up_off = gu_row * NFF, down_off = 2 * up_off; + const size_t blob = down_off + down_row * width; + std::vector inverse; + inverse.reserve(layers * width); + for (const auto& fmt : layout.fmt) { + if (fmt.n_embd != width || fmt.n_ff != NFF || fmt.gu_type != ref::Q3_K || + fmt.d_type != ref::IQ4_NL || fmt.gu_row != gu_row || fmt.d_row != down_row || + fmt.up_off != up_off || fmt.down_off != down_off || fmt.bytes != blob || + !fmt.gu_colscale || fmt.had_gu || fmt.had_d) { + err = "metal latent: expected a uniform colscaled Q3_K/IQ4_NL expert layout"; + return false; + } + for (uint32_t i = 0; i < width; ++i) { + const float v = fmt.gu_colscale[i]; + if (!std::isfinite(v) || v <= 0.0f) { + err = "metal latent: invalid inverse scale"; + return false; + } + inverse.push_back(v); + } + } + if (!latent_dense_.open(dense, layers, groups, N, width, err)) return false; + if (latent_dense_.fmt() != 4) { + err = "metal latent: D/U must be Q4_0 (STRATA_LATENT_DENSE_FMT=4)"; + return false; + } + if (!latent_experts_.open(experts, layers, NEXP, blob, err)) return false; + const char* labels = std::getenv("STRATA_LATENT_LABELS"); + const auto label_path = labels && *labels ? std::filesystem::path(labels) + : std::filesystem::path(experts).parent_path() / "labels.bin"; + std::ifstream lf(label_path, std::ios::binary); + std::vector groups_by_expert(layers * NEXP); + const size_t label_bytes = groups_by_expert.size() * sizeof(int32_t); + if (!lf.read(reinterpret_cast(groups_by_expert.data()), label_bytes) || + lf.peek() != std::char_traits::eof() || + std::any_of(groups_by_expert.begin(), groups_by_expert.end(), + [](int32_t group) { return group < 0 || group >= int32_t(groups); })) { + err = "metal latent: invalid labels at " + label_path.string(); + return false; + } + const size_t layer_du = latent_dense_.d_bytes() + latent_dense_.u_bytes(); + const size_t layer_a = NEXP * blob; + latent_du_ = cx_->wrap(const_cast(latent_dense_.d(0)), latent_dense_.bytes(), layer_du); + latent_a_ = cx_->wrap(const_cast(latent_experts_.blob(0, 0)), layers * layer_a, layer_a); + if (!latent_du_.handle || !latent_a_.handle) { + err = "metal latent: cannot wrap the D/U or expert file"; + return false; + } + latent_labels_ = cx_->alloc(label_bytes); + std::memcpy(latent_labels_.ptr, groups_by_expert.data(), label_bytes); + latent_inv_ = cx_->alloc(inverse.size() * sizeof(float)); + std::memcpy(latent_inv_.ptr, inverse.data(), latent_inv_.bytes); + latent_d_ = cx_->alloc(sizeof(float) * MAXT * groups * width); + latent_down_ = cx_->alloc(sizeof(float) * MAXT * TOPK * width); + const char* grouped = std::getenv("STRATA_METAL_LATENT_GROUP_U"); + latent_grouped_u_ = grouped && !std::strcmp(grouped, "1"); + if (latent_grouped_u_) latent_group_active_ = cx_->alloc(sizeof(uint32_t) * MAXT * groups); + latent_args_ = {0, groups, N, width, NFF, TOPK, P_TOKENT, uint32_t(latent_dense_.u_row_bytes()), + uint32_t(gu_row), uint32_t(down_row), uint32_t(up_off), uint32_t(down_off), uint32_t(blob)}; + const char* verify_fast = std::getenv("STRATA_METAL_LATENT_VERIFY_FAST"); + latent_fast_verify_ = latent_grouped_u_ && verify_fast && std::strcmp(verify_fast, "1") == 0; + latent_lib_ = cx_->compile({"quant.metal", "units.metal", "latent_routed.metal"}); + latent_on_ = true; + std::fprintf(stderr, "strata metal: l%u Q3_K/IQ4_NL experts and Q4_0 D/U ready\n", width); + if (latent_grouped_u_) std::fprintf(stderr, "strata metal: pool routed rows by group before U\n"); + return true; +} + +void Window::latent_routed(std::vector& ds, int layer, int ntok) { + const auto& base = latent_args_; + auto args = base; + args.n_tok = ntok; + const size_t du_offset = layer * (latent_dense_.d_bytes() + latent_dense_.u_bytes()); + const size_t a_offset = layer * size_t(NEXP) * args.expert_bytes; + const size_t label_offset = layer * NEXP * sizeof(int32_t); + const size_t scale_offset = layer * args.latent * sizeof(float); + const uint32_t entries = ntok * args.topk; + static const char* decode_tiles = std::getenv("STRATA_METAL_LATENT_DECODE_TILES"); + const bool narrow = ntok == 1 && decode_tiles && std::strcmp(decode_tiles, "1") == 0; + const bool verify = fast_verify(ntok); + const uint32_t gate_rows = verify ? 32 : narrow ? 16 : 4; + const gguf::Tensor d_weight{"", {N, args.groups * args.latent}, ref::Q4_0, + du_offset, latent_dense_.d_bytes()}; + gemv(ds, d_weight, XF_, latent_d_, args.groups * args.latent, ntok, 0, &latent_du_); + auto run = [&](const char* name, std::vector buffers, uint32_t grid, uint32_t rows) { + mtl::Dispatch d; + d.p = P(latent_lib_, name, {}, name); + d.args = std::move(buffers); + d.args.push_back(mtl::Arg::V(args)); + d.grid[0] = grid; + d.grid[1] = rows; + d.tg[0] = 128; + push(ds, d); + }; + run("latent_colscale", {mtl::Arg::B(latent_d_), mtl::Arg::B(latent_inv_, scale_offset)}, + (ntok * args.groups * args.latent + 127) / 128, 1); + const char* gate_kernel = verify ? "latent_gate_up_4" + : narrow ? "latent_gate_up_8" : "latent_gate_up"; + run(gate_kernel, + {mtl::Arg::B(latent_a_, a_offset), + mtl::Arg::B(latent_labels_, label_offset), mtl::Arg::B(ids_), mtl::Arg::B(latent_d_), mtl::Arg::B(h_)}, + args.ff / gate_rows, entries); + const bool down_narrow = verify || narrow; + run(down_narrow ? "latent_down_4" : "latent_down", {mtl::Arg::B(latent_a_, a_offset), mtl::Arg::B(ids_), + mtl::Arg::B(h_), mtl::Arg::B(latent_down_)}, args.latent / (down_narrow ? 32 : 4), entries); + const auto u_weights = mtl::Arg::B(latent_du_, du_offset + latent_dense_.d_bytes()); + if (latent_grouped_u_) { + const uint32_t group_rows = ntok * args.groups; + // D activations are no longer live after gate/up. Reuse that buffer for the grouped rows. + run("latent_group_reduce", {mtl::Arg::B(latent_labels_, label_offset), mtl::Arg::B(ids_), + mtl::Arg::B(wts_), mtl::Arg::B(latent_down_), mtl::Arg::B(latent_d_), + mtl::Arg::B(latent_group_active_)}, (group_rows * args.latent + 127) / 128, 1); + const char* u_kernel = verify ? "latent_group_u_4" : narrow ? "latent_group_u_8" : "latent_group_u"; + run(u_kernel, {u_weights, mtl::Arg::B(latent_d_), + mtl::Arg::B(latent_group_active_), mtl::Arg::B(ye_)}, args.hidden / gate_rows, group_rows); + } else { + run("latent_u", {u_weights, mtl::Arg::B(latent_labels_, label_offset), + mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::B(latent_down_), + mtl::Arg::B(ye_)}, (args.hidden + 3) / 4, entries); + } +} + +void Window::gemv(std::vector& ds, const gguf::Tensor& t, const Act& x, const mtl::Buffer& y, + size_t y_stride, int ntok, size_t y_off, const mtl::Buffer* wb) { + const int n_in = (int) t.shape[0], n_out = t.shape.size() > 1 ? (int) t.shape[1] : 1; + const int ty = t.type; + int L = 16, Rr = 1, S = 4, K = 1; + if (n_in >= 6144) { L = 32; Rr = 1; S = 8; K = 2; } + else if (n_in == 640) { L = 8; Rr = 1; S = 4; K = 1; } + else if (n_out <= 640 || ty == ref::Q2_0 || ty == ref::BF16) { L = 32; Rr = 1; S = 4; K = 1; } + static const bool gyro_tiles = !(std::getenv("STRATA_METAL_GYRO_TILES") && std::string(std::getenv("STRATA_METAL_GYRO_TILES")) == "0"); // off: the A/B arm + if (gyro_dense_ && gyro_tiles && ntok >= 2) { + // Gyro's dense weights are Q5_K / Q6_K / Q8_0 (3.6 GB a token). gemv_parity --one on the M4, 2 tokens: Q5_K 2560->10240 + // 261 us with 16,1,4,1 against 210 with 2,1,4,1 (-20%), 2560->6144 158 -> 129, 6144->2560 162 -> 139, 2560->12288 + // 313 -> 247, and the shared down Q8_0 640->2560 36 -> 19 us with 4,2,4,1. (The tile changes the order of the lane sums.) + if (ty == ref::Q5_K && n_in == 2560 && n_out >= 6144) { L = 2; Rr = 1; S = 4; K = 1; } + else if (ty == ref::Q5_K && n_in == 6144) { L = 2; Rr = 1; S = 4; K = 1; } + else if (ty == ref::Q8_0 && n_in == 640) { L = 4; Rr = 2; S = 4; K = 1; } + } + { + // windows of several tokens: NR rows a lane share each x load (the window's cost is its x traffic). The sums + // run in the same order for any NR, so the answers do not change. STRATA_METAL_GEMV_NR (1, 2): the tile. + static const int nr_env = [] { const char* e = std::getenv("STRATA_METAL_GEMV_NR"); return e ? std::atoi(e) : 1; }(); + if (ntok > 1 && (nr_env == 2 || fast_verify(ntok)) && n_out >= 2048) { + if (S == 8 && K == 2) { S = 4; K = 2; } + if (S == 8 && K == 1) S = 4; + if (K == 1 || (L == 32 && K == 2)) Rr = 2; + } + } + if (fast_verify(ntok)) { + if (wb == &latent_du_ && ty == ref::Q4_0) { + L = 8; Rr = 1; S = 4; K = 1; + } else if (!wb && ty == ref::Q8_0) { + if (n_in == N && n_out == n_vocab_) { L = 1; Rr = 1; S = 2; K = 1; } + else if (n_in == 6144 && n_out == N) { L = 16; Rr = 2; S = 4; K = 1; } + else if (n_in == N && n_out >= 10240) { L = 8; Rr = 1; S = 8; K = 1; } + else if (n_in == N && n_out == 6144) { L = 8; Rr = 2; S = 4; K = 1; } + } + } + const std::string fn = "gemv_" + std::to_string(ty) + "_" + std::to_string(L) + "_" + std::to_string(Rr) + "_" + + std::to_string(S) + "_" + std::to_string(K); + const mtl::Pipeline* p = P(lib_, fn, mtl::Constants().i(0, ntok), fn + "#" + std::to_string(ntok)); + GemvArgs a{(uint32_t) n_in, (uint32_t) n_out, (uint32_t) ntok, (uint32_t) ref::row_bytes(t.type, n_in), (uint32_t) n_in, + (uint32_t) y_stride}; + mtl::Dispatch d; + d.p = p; + d.args = {mtl::Arg::B(wb ? *wb : W_, t.offset), mtl::Arg::B(x.b, x.plane_bytes() * xkind(ty)), + mtl::Arg::B(x.b, x.plane_bytes() * 4), mtl::Arg::B(y, y_off), mtl::Arg::V(a)}; + const int rpt = (S / K) * (32 / L) * Rr; + d.grid[0] = (n_out + rpt - 1) / rpt; + d.tg[0] = S * 32; + push(ds, d); +} + +void Window::xprep(std::vector& ds, const mtl::Buffer& x, const Act& out, int n) { + mtl::Dispatch d; + d.p = P(lib_, "xprep", {}, "xprep"); + const uint32_t nn = n, nt = MAXT; + d.args = {mtl::Arg::B(x), mtl::Arg::B(out.b), mtl::Arg::V(nn), mtl::Arg::V(nt)}; + d.tg[0] = 256; + d.grid[0] = (n / 16 * MAXT + 255) / 256; + push(ds, d); +} + +// the hyper-connection read (gr.metal, 3 dispatches): the previous sublayer's write folded in when `apply` +void Window::gr_read(std::vector& ds, const std::string& pre, bool apply, const Act& out, bool final_, + int ntok) { + const gguf::Tensor& tn = T(pre + "_norm.weight"); + const gguf::Tensor& tdn = T(pre + "_down.weight"); + const gguf::Tensor& tup = T(pre + "_up.weight"); + const gguf::Tensor& tin = final_ ? tdn : T(pre + "_inject.weight"); + GrArgs a{N, HC, LR, (uint32_t) ntok, apply ? 1u : 0u, eps_}; + mtl::Dispatch d1; + d1.p = P(lib_, "gr_norm", {}, "gr_norm"); + d1.args = {mtl::Arg::B(R_), mtl::Arg::B(bo_), mtl::Arg::B(inj_[cur_inj_]), mtl::Arg::B(W_, tn.offset), mtl::Arg::B(xn_), + mtl::Arg::V(a)}; + d1.grid[0] = HC; + d1.grid[1] = ntok; + d1.tg[0] = 256; + push(ds, d1); + // each matrix: its load-time Q8_0 copy (q8off_), else the file's own tensor (Unsloth's are Q8_0 already, their + // inject rows F32) + auto src = [&](const gguf::Tensor& t) -> std::pair { + auto it = q8off_.find(t.name); + return it != q8off_.end() ? std::make_pair(&Q8_, it->second) : std::make_pair(&W_, (size_t) t.offset); + }; + auto is_q8 = [&](const gguf::Tensor& t) { return q8off_.count(t.name) || t.type == ref::Q8_0; }; + const auto [bd, od] = src(tdn); + const auto [bi, oi] = src(tin); + const auto [bu, ou] = src(tup); + const bool q8 = is_q8(tdn) && is_q8(tup); + const bool injf = q8 && !final_ && !q8off_.count(tin.name) && tin.type == ref::F32; + mtl::Dispatch d2; + const char* dn = injf ? "gr_down_q8f_8_8" : q8 ? "gr_down_q8_8_8" : "gr_down_8_8"; + d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), std::string(dn) + "#" + std::to_string(ntok)); + d2.args = {mtl::Arg::B(*bd, od), mtl::Arg::B(*bi, oi), mtl::Arg::B(xn_), mtl::Arg::B(lo_), + mtl::Arg::B(inj_[1 - cur_inj_]), mtl::Arg::V(a)}; + d2.grid[0] = LR + HC; + d2.tg[0] = 256; + push(ds, d2); + mtl::Dispatch d3; + const char* up = q8 ? "gr_up_q8" : "gr_up"; + d3.p = P(lib_, up, mtl::Constants().i(1, ntok), std::string(up) + "#" + std::to_string(ntok)); + GrArgs au = a; + au.n_tok = MAXT; // the plane stride of the prepared activation + d3.args = {mtl::Arg::B(*bu, ou), mtl::Arg::B(lo_), mtl::Arg::B(xn_), mtl::Arg::B(out.b), mtl::Arg::V(au)}; + d3.grid[0] = N / 16; + d3.tg[0] = 256; + push(ds, d3); + cur_inj_ = 1 - cur_inj_; +} + +namespace { +struct HadArgs { uint32_t width, x_stride, y_stride; float scale; uint32_t has_signs, pad; }; // quant.metal +} +// y = H (s * x) per block of 128 values of `rows` rows of `width` (x_off / y_off: byte offsets; f16: half rows, no signs) +void Window::had_rows(std::vector& ds, int lib, const std::string& pfx, const mtl::Buffer& x, size_t x_off, + const mtl::Buffer& y, size_t y_off, int rows, int width, const mtl::Buffer* signs, bool f16) { + HadArgs a{(uint32_t) width, (uint32_t) width, (uint32_t) width, 1.0f / std::sqrt((float) had_block_), signs ? 1u : 0u, 0u}; + const char* fn = f16 ? "had128_f16" : "had128_f32"; + mtl::Dispatch d; + d.p = P(lib, fn, {}, pfx + fn); + d.args = {mtl::Arg::B(x, x_off), mtl::Arg::B(y, y_off), mtl::Arg::B(signs ? *signs : x), mtl::Arg::V(a)}; + d.grid[0] = width / 128; + d.grid[1] = rows; + d.tg[0] = 128; + push(ds, d); +} + +void Window::had_rows_ind(std::vector& ds, int lib, const std::string& pfx, const mtl::Buffer& h, + const mtl::Buffer& meta) { + HadArgs a{(uint32_t) NFF, (uint32_t) NFF, (uint32_t) NFF, 1.0f / std::sqrt((float) had_block_), hsh_.handle ? 1u : 0u, 0u}; + mtl::Dispatch d; + d.p = P(lib, "had128_f16_rows", {}, pfx + "had128_f16_rows"); + d.args = {mtl::Arg::B(h), mtl::Arg::B(h), mtl::Arg::B(hsh_.handle ? hsh_ : h), mtl::Arg::V(a)}; + d.tg[0] = 128; + d.indirect = &meta; + d.ind_off = 8 * 4; // (rows, 1, 1): the padded rows of the tiles + push(ds, d); +} + +void Window::layer(std::vector& ds, int l, int ntok, int64_t pos0) { + const std::string b = "blk." + std::to_string(l) + "."; + // ---- the mixer: layer 1's read after the PLE is not applied again (its pending write was applied before it) + gr_read(ds, b + "hc_attn", l > 0 && !(l == 1 && ple_on_), XA_, false, ntok); + if (l % 4 != 3) { + const int gi = l - l / 4; + // a one-token window commits in the step; a longer one leaves the state and commit() folds the kept prefix + // from the inputs kept per layer (qkvL_ ...) + const bool self_commit = ntok == 1; + const uint32_t cm = self_commit ? 1u : !partial_ ? 0u : win_head_ ? (0x100u | 1u) : 1u; + const size_t o_qkv = (size_t) gi * MAXT * 10240 * 4, o_ab = (size_t) gi * MAXT * 48 * 4, o_z = (size_t) gi * MAXT * 6144 * 4; + gemv(ds, T(b + "attn_qkv.weight"), XA_, qkvL_, 10240, ntok, o_qkv); + nb_ = true; + gemv(ds, T(b + "attn_gate.weight"), XA_, zL_, 6144, ntok, o_z); + nb_ = true; + gemv(ds, T(b + "ssm_alpha.weight"), XA_, alL_, 48, ntok, o_ab); + nb_ = true; + gemv(ds, T(b + "ssm_beta.weight"), XA_, beL_, 48, ntok, o_ab); + GdnArgs a{16, 48, 10240, cm, eps_, fuse_ ? (uint32_t) (MAXT * 6144) : 0u}; + mtl::Dispatch d; + d.p = P(lgdn_, "gdn_step", mtl::Constants().i(2, ntok), "gdn_step#" + std::to_string(ntok)); + const int hp = hp_; + if (rows_) { // each row is a different sequence: its own state, parity and conv history (rinfo_), one grid row each + d.p = P(lgdn_, "gdn_step", mtl::Constants().i(2, 1).b(31, true), "gdn_step#rows"); + GdnArgs ar{16, 48, 10240, 1u, eps_, fuse_ ? (uint32_t) (MAXT * 6144) : 0u}; + d.args = {mtl::Arg::B(gstBig_[(size_t) gi]), mtl::Arg::B(hbBig_[(size_t) gi]), mtl::Arg::B(qkvL_, o_qkv), + mtl::Arg::B(alL_, o_ab), mtl::Arg::B(beL_, o_ab), mtl::Arg::B(zL_, o_z), + mtl::Arg::B(W_, T(b + "ssm_conv1d.weight").offset), + mtl::Arg::B(W_, T(b + "ssm_dt.bias").offset), mtl::Arg::B(W_, T(b + "ssm_a").offset), + mtl::Arg::B(W_, T(b + "ssm_norm.weight").offset), mtl::Arg::B(fuse_ ? XY_.b : yg_), mtl::Arg::V(ar), + mtl::Arg::B(hbBig_[(size_t) gi]), mtl::Arg::B(rinfo_)}; + d.grid[0] = 48; + d.grid[1] = (uint32_t) rows_n_; + d.tg[0] = 512; + push(ds, d); + if (!fuse_) xprep(ds, yg_, XY_, 6144); + gemv(ds, T(b + "ssm_out.weight"), XY_, bo_, N, ntok); + } else { + d.args = {mtl::Arg::B(gst_[(size_t) gi]), mtl::Arg::B(ghist_[hp][(size_t) gi]), mtl::Arg::B(qkvL_, o_qkv), + mtl::Arg::B(alL_, o_ab), mtl::Arg::B(beL_, o_ab), mtl::Arg::B(zL_, o_z), + mtl::Arg::B(W_, T(b + "ssm_conv1d.weight").offset), + mtl::Arg::B(W_, T(b + "ssm_dt.bias").offset), mtl::Arg::B(W_, T(b + "ssm_a").offset), + mtl::Arg::B(W_, T(b + "ssm_norm.weight").offset), mtl::Arg::B(fuse_ ? XY_.b : yg_), mtl::Arg::V(a), + mtl::Arg::B(ghist_[1 - hp][(size_t) gi])}; + d.grid[0] = 48; + d.tg[0] = 512; + push(ds, d); + (void) self_commit; // the history ping-pong turns after the window (run) or at commit + if (!fuse_) xprep(ds, yg_, XY_, 6144); + gemv(ds, T(b + "ssm_out.weight"), XY_, bo_, N, ntok); + } + } else { + const int qi = l / 4; + gemv(ds, T(b + "attn_q.weight"), XA_, qfull_, 12288, ntok); + nb_ = true; + gemv(ds, T(b + "attn_k.weight"), XA_, kcur_, 512, ntok); + nb_ = true; + gemv(ds, T(b + "attn_v.weight"), XA_, vcur_, 512, ntok); + nb_ = true; + gemv(ds, T(b + "indexer.q_proj.weight"), XA_, iqr_, 512, ntok); + nb_ = true; + gemv(ds, T(b + "indexer.k_proj.weight"), XA_, ikr_, 128, ntok); + QsaArgs a{(uint32_t) pos0, 0, (uint32_t) max_cells_, (uint32_t) ntok, eps_}; + mtl::Dispatch d; + const bool rw = rows_; // (batch slots: pos0 is the highest row's cell, the rows have their own in rinfo_) + const mtl::Constants rc = mtl::Constants().b(30, true); + const mtl::Buffer& KQ = rw ? kqBig_[(size_t) qi] : kq_[(size_t) qi]; + const mtl::Buffer& VQ = rw ? vqBig_[(size_t) qi] : vq_[(size_t) qi]; + const mtl::Buffer& KS = rw ? kscBig_[(size_t) qi] : ksc_[(size_t) qi]; + const mtl::Buffer& VS = rw ? vscBig_[(size_t) qi] : vsc_[(size_t) qi]; + const mtl::Buffer& TL = rw ? tailBig_[(size_t) qi] : tail_[(size_t) qi]; + const mtl::Buffer& DD = rw ? deadBig_[(size_t) qi] : dead_[(size_t) qi]; + const mtl::Buffer& PL = rw ? pooledBig_[(size_t) qi] : pooled_[(size_t) qi]; + d.p = rw ? P(lqsa_, "qsa_prep", rc, "qsa_prep#rows") : P(lqsa_, "qsa_prep", qsa_c(), qsa_k()); + d.args = {mtl::Arg::B(qfull_), mtl::Arg::B(kcur_), mtl::Arg::B(vcur_), mtl::Arg::B(iqr_), mtl::Arg::B(ikr_), + mtl::Arg::B(W_, T(b + "attn_q_norm.weight").offset), mtl::Arg::B(W_, T(b + "attn_k_norm.weight").offset), + mtl::Arg::B(W_, T(b + "indexer.q_norm.weight").offset), mtl::Arg::B(W_, T(b + "indexer.k_norm.weight").offset), + mtl::Arg::B(cs_), mtl::Arg::B(qcur_), mtl::Arg::B(qidx_), mtl::Arg::B(KQ), + mtl::Arg::B(VQ), mtl::Arg::B(KS), mtl::Arg::B(VS), + mtl::Arg::B(TL), mtl::Arg::B(DD), mtl::Arg::B(PL), + mtl::Arg::V(a), mtl::Arg::B(mrope_)}; + if (rw) d.args.push_back(mtl::Arg::B(rinfo_)); + d.grid[0] = 33; + d.grid[1] = ntok; + d.tg[0] = 256; + push(ds, d); + if (rec_) rec_->push_back({(int) ds.size() - 1, 0}); + mtl::Dispatch s; + s.p = rw ? P(lqsa_, "qsa_score", rc, "qsa_score#rows") : P(lqsa_, "qsa_score", {}, "qsa_score"); + s.args = {mtl::Arg::B(qidx_), mtl::Arg::B(PL), mtl::Arg::B(DD), mtl::Arg::B(scores_), + mtl::Arg::V(a)}; + if (rw) s.args.push_back(mtl::Arg::B(rinfo_)); + s.grid[0] = (uint32_t) (((pos0 + ntok) / 4 + 1 + 7) / 8); + s.grid[1] = ntok; + s.tg[0] = 256; + const bool sel = !(fuse_ && pos0 + ntok <= 2051); + if (sel) { + push(ds, s); + if (rec_) rec_->push_back({(int) ds.size() - 1, 2}); + } + mtl::Dispatch k; + k.p = rw ? P(lqsa_, "qsa_topk2", rc, "qsa_topk#rows") : P(lqsa_, "qsa_topk2", {}, "qsa_topk"); + k.args = {mtl::Arg::B(scores_), mtl::Arg::B(qids_), mtl::Arg::V(a)}; + if (rw) k.args.push_back(mtl::Arg::B(rinfo_)); + k.grid[0] = ntok; + k.tg[0] = 1024; + if (sel) { + push(ds, k); + if (rec_) rec_->push_back({(int) ds.size() - 1, 0}); + } + mtl::Dispatch at; + at.p = rw ? P(lqsa_, "qsa_attn", rc, "qsa_attn#rows") : P(lqsa_, "qsa_attn", {}, "qsa_attn"); + at.args = {mtl::Arg::B(qcur_), mtl::Arg::B(KQ), mtl::Arg::B(VQ), mtl::Arg::B(KS), + mtl::Arg::B(VS), mtl::Arg::B(qids_), mtl::Arg::B(part_), mtl::Arg::V(a)}; + if (rw) at.args.push_back(mtl::Arg::B(rinfo_)); + at.grid[0] = 33; + at.grid[1] = 2; + at.grid[2] = ntok; + at.tg[0] = 256; + push(ds, at); + if (rec_) rec_->push_back({(int) ds.size() - 1, 0}); + mtl::Dispatch m; + m.p = rw ? P(lqsa_, "qsa_merge", rc, "qsa_merge#rows") : P(lqsa_, "qsa_merge", {}, "qsa_merge"); + QsaArgs am = a; + am.n_tok = MAXT; // the plane stride + m.args = {mtl::Arg::B(part_), mtl::Arg::B(qfull_), mtl::Arg::B(XQ_.b), mtl::Arg::V(am)}; + if (rw) m.args.push_back(mtl::Arg::B(rinfo_)); + m.grid[0] = 24; + m.grid[1] = ntok; + m.tg[0] = 256; + push(ds, m); + if (rec_) rec_->push_back({(int) ds.size() - 1, 1}); + gemv(ds, T(b + "attn_output.weight"), XQ_, bo_, N, ntok); + } + // ---- the MoE + gr_read(ds, b + "hc_ffn", true, XF_, false, ntok); + const gguf::Tensor& trt = T(b + "ffn_gate_inp.weight"); + const bool rf = !latent_on_ && routef_ && (trt.type == ref::BF16 || trt.type == ref::F32); + if (!rf) { + gemv(ds, trt, XF_, logits_, 512, ntok); + nb_ = true; + } + if (stream_ && st_pred_ && l + 1 < 48) { // the next layer's routing guessed from this layer's FFN input + nb_ = !rf; // beside the router GEMV; with router_route it is the first reader of XF_ here: a barrier + gemv(ds, T("blk." + std::to_string(l + 1) + ".ffn_gate_inp.weight"), XF_, predL_, 512, ntok, + (size_t) l * MAXT * 512 * 4); + if (predL2_.ptr && l + 2 < 48) { + nb_ = true; + gemv(ds, T("blk." + std::to_string(l + 2) + ".ffn_gate_inp.weight"), XF_, predL2_, 512, ntok, + (size_t) l * MAXT * 512 * 4); + } + nb_ = true; + } + gemv(ds, T(b + "ffn_gate_shexp.weight"), XF_, gsh_, NFF, ntok); + nb_ = true; + gemv(ds, T(b + "ffn_up_shexp.weight"), XF_, ush_, NFF, ntok); + nb_ = true; + gemv(ds, T(b + "ffn_gate_inp_shexp.weight"), XF_, sg_, 1, ntok); + if (latent_on_) { + const bool overlap = fast_verify(ntok); + MoeArgs a{}; + a.n_tok = ntok; + a.topk = TOPK; + a.n_expert = NEXP; + mtl::Dispatch route; + route.p = overlap ? P(lib_, "moe_route", mtl::Constants().b(28, true), "moe_route#only") + : P(lib_, "moe_route", {}, "moe_route"); + route.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), + mtl::Arg::B(plan_), mtl::Arg::V(a)}; + route.tg[0] = overlap ? 32 * ntok : 512; + push(ds, route); + if (!overlap) latent_routed(ds, l, ntok); + + const GlueArgs glue{NFF, MAXT, N, TOPK}; + mtl::Dispatch shared; + shared.p = P(lib_, "swiglu_prep", {}, "swiglu_prep"); + shared.args = {mtl::Arg::B(gsh_), mtl::Arg::B(ush_), mtl::Arg::B(XH_.b), mtl::Arg::V(glue)}; + shared.grid[0] = NFF / 16; + shared.grid[1] = ntok; + shared.tg[0] = 16; + nb_ = overlap; + push(ds, shared); + if (overlap) { + // Route, shared SwiGLU, and D projection read independent inputs. + nb_ = true; + latent_routed(ds, l, ntok); + nb_ = true; // the shared down projection is independent of the last routed dispatch + } + gemv(ds, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); + } else { + MoeArgs a{}; + const gguf::Tensor &tg_ = T(b + "ffn_gate_exps.weight"), &tu = T(b + "ffn_up_exps.weight"), + &td = T(b + "ffn_down_exps.weight"); + a.gate_off = tg_.offset; + a.up_off = tu.offset; + a.down_off = td.offset; + if (stream_) { // offsets inside the expert's arena slot + a.gate_off = 0; + a.up_off = st_->up_in_slot(l); + a.down_off = st_->down_in_slot(l); + } + const size_t ebase = moe_rebase(a.gate_off, a.up_off, a.down_off); + a.gu_row = (uint32_t) ref::row_bytes(tg_.type, N); + a.d_row = (uint32_t) ref::row_bytes(td.type, NFF); + a.gu_stride = (uint64_t) a.gu_row * NFF; + a.d_stride = (uint64_t) a.d_row * N; + a.n_embd = N; + a.n_ff = NFF; + a.n_tok = ntok; + a.topk = TOPK; + a.n_expert = NEXP; + const bool dq2 = td.type == ref::Q2_0; + // the trellis gate/up: 16 lanes a row, 2 rows a lane (expert_parity --bench --sweep, M4, TQK6, 10 experts: 166.6 us a layer + // against 174.5 for 32,2,4 and 196.8 for 32,1,4; also the fastest at 2 and 3 tokens) + const bool tq_ = tg_.type >= 144; + const int gl = tq_ ? 16 : 32, gr = tq_ ? 2 : 1, gs = 4, ds_ = 4; + // the down tile: Q2_0 8,1; else 4,2 - but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there: + // 24-36 GB/s, mac-kern 513fc615) + const bool m4 = !m5_layout(); + const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4, + dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2; + a.gx_gu = NFF / (gs * (32 / gl) * gr); + a.gx_dn = N / (ds_ * (32 / dl) * dr); + mtl::Dispatch r; + if (rf) { + // the router GEMV + moe_route (bitwise both), beside the shared expert's GEMVs + r.p = P(lib_, "router_route_" + std::to_string(trt.type), mtl::Constants().i(0, ntok), "router_route_" + std::to_string(trt.type) + "#" + std::to_string(ntok)); + GemvArgs ga{(uint32_t) N, 512, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, 512}; + r.args = {mtl::Arg::B(W_, trt.offset), mtl::Arg::B(XF_.b, 0), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4), + mtl::Arg::B(logits_), mtl::Arg::V(ga), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), + mtl::Arg::V(a), mtl::Arg::B(rcounter_)}; + r.grid[0] = 512 / 16; + r.tg[0] = 512; + nb_ = true; + } else { + r.p = P(lib_, "moe_route", {}, "moe_route"); + r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)}; + r.tg[0] = 512; + } + push(ds, r); + { + GlueArgs ga{NFF, MAXT, N, TOPK}; + mtl::Dispatch d; + d.p = P(lib_, "swiglu_prep", {}, "swiglu_prep"); + d.args = {mtl::Arg::B(gsh_), mtl::Arg::B(ush_), mtl::Arg::B(XH_.b), mtl::Arg::V(ga)}; + d.grid[0] = NFF / 16; + d.grid[1] = ntok; + d.tg[0] = 16; + nb_ = !rf; // with router_route the shared GEMVs are not behind a barrier yet: this one waits for them + push(ds, d); + } + const std::string gun = "moe_gu_" + std::to_string(tg_.type) + "_" + std::to_string(gl) + "_" + std::to_string(gr) + + "_" + std::to_string(gs); + const std::string dnn = "moe_down_" + std::to_string(dq2 ? 142 : td.type) + "_" + std::to_string(dl) + "_" + + std::to_string(dr) + "_" + std::to_string(ds_); + mtl::Dispatch g; + const mtl::Buffer& EW = stream_ ? st_->chunk(0) : W_; + // function constant 23 (QUANT_CTAB): the i-quant grids from constant memory, no threadgroup copy (same values) + const bool ct = metal_ctab(); + mtl::Constants ac = stream_ ? mtl::Constants().b(20, true) : mtl::Constants(); + ac.b(23, ct); + const std::string ak = std::string(stream_ ? "#arena" : "") + (ct ? "#ct" : ""); + g.p = P(lib_, gun, ac, gun + ak); + const bool hg = had_gu(l), hd = had_dn(l); + if (hg) had_rows(ds, lib_, "d:", XF_.b, 0, XR_, 0, ntok, N, hsx_.handle ? &hsx_ : nullptr, false); // before the host's turn + g.args = {mtl::Arg::B(EW, ebase), hg ? mtl::Arg::B(XR_) : mtl::Arg::B(XF_.b, XF_.plane_bytes() * xkind(tg_.type)), mtl::Arg::B(h_), + mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4)}; + if (stream_) + for (int c = -1; c < 3; ++c) + g.args.push_back(c < 0 ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(st_->chunk(c + 1))); + g.tg[0] = 32 * gs; + g.indirect = &plan_; + g.ind_off = P_IGU * 4; + if (cuts_) cuts_->push_back((int) ds.size()); // the host's turn: this layer's experts put in place + push(ds, g); + nb_ = true; + gemv(ds, T(b + "ffn_down_shexp.weight"), XH_, ysh_, N, ntok); + if (hd) had_rows(ds, lib_, "d:", h_, 0, h_, 0, ntok * TOPK, NFF, hsh_.handle ? &hsh_ : nullptr, false); // H h, in place + mtl::Dispatch dd; + dd.p = P(lib_, dnn, ac, dnn + ak); + dd.args = {mtl::Arg::B(EW, ebase), mtl::Arg::B(h_), mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(h_)}; + if (stream_) + for (int c = -1; c < 3; ++c) + dd.args.push_back(c < 0 ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(st_->chunk(c + 1))); + dd.tg[0] = 32 * ds_; + dd.indirect = &plan_; + dd.ind_off = P_IDN * 4; + push(ds, dd); + } + { + mtl::Dispatch d; + if (latent_grouped_u_) { + auto a = latent_args_; + a.n_tok = ntok; + d.p = P(latent_lib_, "latent_group_combine", {}, "latent_group_combine"); + d.args = {mtl::Arg::B(ye_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::B(bo_), mtl::Arg::V(a)}; + } else { + GlueArgs a{N, (uint32_t) ntok, N, TOPK}; + d.p = P(lib_, "moe_combine_sh", {}, "moe_combine_sh"); + d.args = {mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::B(bo_), mtl::Arg::V(a)}; + } + d.grid[0] = (ntok * N + 255) / 256; + d.tg[0] = 256; + push(ds, d); + } +} + +namespace { +float bf16f(uint16_t h) { uint32_t i = (uint32_t) h << 16; float f; std::memcpy(&f, &i, 4); return f; } +// ple.cu's gnorm_kernel: per stream, the squares rounded in f32, summed in double +void gnorm(const float* x, const float* w, float* y, int n_embd, int hc, float eps) { + for (int c = 0; c < hc; ++c) { + const float* xc = x + (size_t) c * n_embd; + double acc = 0.0; + for (int d = 0; d < n_embd; ++d) { const float sq = xc[d] * xc[d]; acc += (double) sq; } + const float mean = (float) (acc / (double) n_embd); + const float scale = 1.0f / std::sqrt(mean + eps); + for (int d = 0; d < n_embd; ++d) y[(size_t) c * n_embd + d] = xc[d] * scale * w[(size_t) c * n_embd + d]; + } +} +} // namespace + +// The PLE block (ple.hpp) for the window's T rows on the CPU, after layer 0 (its pending write applied to R): R += +// gated + silu(conv); the normalized rows go to last_norm_ (the history advances at commit). +void Window::ple_cpu(int Tn) { + using namespace strata::kernels; + const int hc_dim = NG_HC_DIM; + const float* nkey = (const float*) ((const uint8_t*) W_.ptr + T("blk.1.ple_norm_key.weight").offset); + const float* nqry = (const float*) ((const uint8_t*) W_.ptr + T("blk.1.ple_norm_query.weight").offset); + const float* nconv = (const float*) ((const uint8_t*) W_.ptr + T("blk.1.ple_norm_conv.weight").offset); + const uint16_t* kw = kw_host(); + std::vector key(hc_dim), query(hc_dim), gated(hc_dim), norm(hc_dim), hist = ple_hist_; + for (int t = 0; t < Tn; ++t) { + float* hidden = R_.as() + (size_t) t * HC * N; + gnorm(pkey_.as() + (size_t) t * hc_dim, nkey, key.data(), N, HC, NG_RMS_EPS); + gnorm(hidden, nqry, query.data(), N, HC, NG_RMS_EPS); + const float* value = pval_.as() + (size_t) t * N; + float gate[HC]; + const float inv_sqrt_n = 1.0f / std::sqrt((float) N); + for (int c = 0; c < HC; ++c) { + double acc = 0.0; + for (int d = 0; d < N; ++d) acc += (double) (key[(size_t) c * N + d] * query[(size_t) c * N + d]); + const float s = (float) acc * inv_sqrt_n; + const float mag = std::sqrt(std::max(std::fabs(s), 1e-6f)); + const float sgn = s > 0.0f ? 1.0f : (s < 0.0f ? -1.0f : 0.0f); + gate[c] = 1.0f / (1.0f + std::exp(-(sgn * mag))); + } + for (int i = 0; i < hc_dim; ++i) gated[(size_t) i] = value[i % N] * gate[i / N]; + gnorm(gated.data(), nconv, norm.data(), N, HC, NG_RMS_EPS); + for (int c = 0; c < hc_dim; ++c) { + float acc = 0.0f; + for (int k = 0; k < PLE_CONV_KERNEL; ++k) { + const int row = NG_HIST - (PLE_CONV_KERNEL - 1 - k) * NGRAM_SIZE; + const float v = row == NG_HIST ? norm[(size_t) c] : hist[(size_t) row + (size_t) NG_HIST * c]; + acc += ref::f16_to_f32(kw[k + PLE_CONV_KERNEL * c]) * v; + } + const float conv = acc / (1.0f + std::exp(-acc)); + hidden[c] = hidden[c] + gated[(size_t) c] + conv; + } + std::memcpy(last_norm_.data() + (size_t) t * hc_dim, norm.data(), (size_t) hc_dim * 4); + // the next row of this window sees this one in its history + for (int c = 0; c < hc_dim; ++c) { + float* col = hist.data() + (size_t) c * NG_HIST; + for (int r = 0; r + 1 < NG_HIST; ++r) col[r] = col[r + 1]; + col[NG_HIST - 1] = norm[(size_t) c]; + } + } +} + +bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, std::string& err) { + mtp_src_ = nullptr; + if (Tn < 1 || Tn > MAXT || (head && Tn > max_t_)) { err = "metal window: window size out of range"; return false; } + if (pos0 + Tn + 1 >= max_cells_) { err = "metal window: the window runs past the context"; return false; } + const double t0 = now_ms(); + // ---- the embeddings (CPU dequant of the token rows straight into the shared buffer), broadcast to the streams + const gguf::Tensor& te = T("token_embd.weight"); + const size_t rb = ref::row_bytes(te.type, N); + for (int t = 0; t < Tn; ++t) { + float* r = R_.as() + (size_t) t * HC * N; + if (const float* img = img_row(pos0 + t)) std::memcpy(r, img, (size_t) N * 4); + else ref::dequant(te.type, (const uint8_t*) W_.ptr + te.offset + (size_t) tokens[t] * rb, r, N); + for (int c = 1; c < HC; ++c) std::memcpy(r + (size_t) c * N, r, (size_t) N * 4); + } + // ---- the PLE rows (the table read is the shared host code's), projected on the GPU + if (ple_on_) { + strata::core::SessionState& ss = *ss_; + uint32_t rows[MAXT * strata::kernels::PLE_N_HEADS]; + int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]}; + for (int t = 0; t < Tn; ++t) { + if (rows_) { prev[0] = rows_ss_[t]->ple_prev[0]; prev[1] = rows_ss_[t]->ple_prev[1]; } // (batch slots: each row its own sequence) + strata::kernels::ngram_rows(&tokens[t], prev, 1, main_ss_->ple.consts, rows + t * strata::kernels::PLE_N_HEADS); + if (!rows_) { prev[0] = prev[1]; prev[1] = tokens[t]; } + } + if (!main_ss_->ple.table->gather_batch(rows, (size_t) Tn, emb_.as(), err)) return false; + } + tp_pre_ = now_ms(); + const bool sel = !(fuse_ && pos0 + Tn <= 2051); + if (sel) qdirty_ = true; + else if (qdirty_) { + for (int t = 0; t < MAXT; ++t) + for (int c = 0; c < 2051; ++c) qids_.as()[t * 2051 + c] = c; + qdirty_ = false; + } + const int key = rows_ ? 1000000 + Tn * 2 + (sel ? 1 : 0) + (int) (pend_fold_ * 100) + : (((((Tn * 2 + (head ? 1 : 0)) * 2 + hp_) * (MAXT + 1)) + pend_fold_) * 2 + (sel ? 1 : 0)) * 2 + + (mrope_on() ? 1 : 0); + mrope_sync(pos0, Tn); + const int fold_n = pend_fold_; + pend_fold_ = 0; + win_head_ = head; + auto hit = wcache_.find(key); + if (hit != wcache_.end()) { + // the same window shape: only the positions change (QSA's cells, the score grid) + Cached& c = hit->second; + for (const auto& [i, kind] : c.qsa) { + mtl::Dispatch& d = c.ds[(size_t) i]; + QsaArgs a{(uint32_t) pos0, 0, (uint32_t) max_cells_, (uint32_t) (kind == 1 ? MAXT : Tn), eps_}; + // the QsaArgs: the dispatch's last inline argument (qsa_prep has the M-RoPE table bound after it) + mtl::Arg* qa = nullptr; + for (auto& g : d.args) + if (!g.bytes.empty()) qa = &g; + std::memcpy(qa->bytes.data(), &a, sizeof a); + if (kind == 2) d.grid[0] = (uint32_t) (((pos0 + Tn) / 4 + 1 + 7) / 8); + } + seg_T_ = Tn; + const double r0 = now_ms(); + last_gpu_ms_ = 1000.0 * run_ds(c); + { + static const int kt = std::getenv("STRATA_METAL_KPROF") ? std::atoi(std::getenv("STRATA_METAL_KPROF")) : 0; + static int seen = 0; + if (kt == Tn && head && ++seen == 20 && !stream_) kprof(c.ds, Tn); + } + tp_rg_ = now_ms() - r0; + last_n_disp_ = (int) c.ds.size(); + return rows_ ? finish_rows(Tn, tokens, out, t0) : finish(Tn, tokens, pos0, head, out, t0); + } + if (stream_ && !st_->ready() && !stream_arena(err)) return false; + Cached& cw = wcache_[key]; + rec_ = &cw.qsa; + cuts_ = stream_ ? &cw.cuts : nullptr; + std::vector& ds = cw.ds; + nb_ = false; + if (fold_n > 0) fold(ds, fold_n, 1 - hp_); // the last window's commit (hp_ already turned) + if (ple_on_) { + xprep(ds, emb_, XE_, N); + gemv(ds, T("blk.1.ple_key.weight"), XE_, pkey_, 10240, Tn); + nb_ = true; + gemv(ds, T("blk.1.ple_value.weight"), XE_, pval_, N, Tn); + } + cur_inj_ = 0; + layer(ds, 0, Tn, pos0); + if (ple_on_) { + // layer 0's FFN write applied to R (gr_norm with apply writes R back; its xn is not used) + const gguf::Tensor& tn = T("blk.1.hc_attn_norm.weight"); + GrArgs a{N, HC, LR, (uint32_t) Tn, 1u, eps_}; + mtl::Dispatch d1; + d1.p = P(lib_, "gr_norm", {}, "gr_norm"); + d1.args = {mtl::Arg::B(R_), mtl::Arg::B(bo_), mtl::Arg::B(inj_[cur_inj_]), mtl::Arg::B(W_, tn.offset), mtl::Arg::B(xn_), + mtl::Arg::V(a)}; + d1.grid[0] = HC; + d1.grid[1] = Tn; + d1.tg[0] = 256; + push(ds, d1); + // the PLE block on the GPU (msl/ple.metal): norms, gate, conv, the residual update + struct PleArgs { uint32_t n_embd, hc; float eps; } pa{N, HC, strata::kernels::NG_RMS_EPS}; + auto off = [&](const char* n) { return T(n).offset; }; + mtl::Dispatch g1, g2, g3, g4; + g1.p = P(lple_, "ple_gnorm", {}, "ple_gnorm"); + g1.args = {mtl::Arg::B(pkey_), mtl::Arg::B(W_, off("blk.1.ple_norm_key.weight")), mtl::Arg::B(keyn_), mtl::Arg::V(pa)}; + g1.grid[0] = HC; g1.grid[1] = Tn; g1.tg[0] = 256; + push(ds, g1); + g2.p = g1.p; + g2.args = {mtl::Arg::B(R_), mtl::Arg::B(W_, off("blk.1.ple_norm_query.weight")), mtl::Arg::B(qn_), mtl::Arg::V(pa)}; + g2.grid[0] = HC; g2.grid[1] = Tn; g2.tg[0] = 256; + nb_ = true; // beside the key norm + push(ds, g2); + g3.p = P(lple_, "ple_gate_norm", {}, "ple_gate_norm"); + g3.args = {mtl::Arg::B(keyn_), mtl::Arg::B(qn_), mtl::Arg::B(pval_), mtl::Arg::B(W_, off("blk.1.ple_norm_conv.weight")), + mtl::Arg::B(gated_), mtl::Arg::B(normW_), mtl::Arg::V(pa)}; + g3.grid[0] = HC; g3.grid[1] = Tn; g3.tg[0] = 256; + push(ds, g3); + g4.p = rows_ ? P(lple_, "ple_conv_add", mtl::Constants().b(32, true), "ple_conv_add#rows") : P(lple_, "ple_conv_add", {}, "ple_conv_add"); + g4.args = {mtl::Arg::B(rows_ ? pleBig_ : pleHist_), mtl::Arg::B(normW_), kw_arg(), + mtl::Arg::B(gated_), mtl::Arg::B(R_), mtl::Arg::V(pa)}; + if (rows_) g4.args.push_back(mtl::Arg::B(rinfo_)); + g4.grid[0] = HC * N / 256; g4.grid[1] = Tn; g4.tg[0] = 256; + push(ds, g4); + } + for (int l = 1; l < 48; ++l) layer(ds, l, Tn, pos0); + // ---- the final mixer and the head + gr_read(ds, "output_hc", true, XA_, true, Tn); + if (head) gemv(ds, T("output.weight"), XA_, lm_, (size_t) n_vocab_, Tn); + if (head) { + ArgArgs a{(uint32_t) n_vocab_, 120}; + mtl::Dispatch d1, d2; + d1.p = P(lsmp_, "argmax_part", {}, "argmax_part"); + d1.args = {mtl::Arg::B(lm_), mtl::Arg::B(apv_), mtl::Arg::B(api_), mtl::Arg::V(a)}; + d1.grid[0] = 120; + d1.grid[1] = Tn; + d1.tg[0] = 256; + push(ds, d1); + d2.p = P(lsmp_, "argmax_final", {}, "argmax_final"); + d2.args = {mtl::Arg::B(apv_), mtl::Arg::B(api_), mtl::Arg::B(tokout_), mtl::Arg::V(a)}; + d2.grid[0] = Tn; + d2.tg[0] = 256; + push(ds, d2); + } + rec_ = nullptr; + cuts_ = nullptr; + seg_T_ = Tn; + const double r0 = now_ms(); + last_gpu_ms_ = 1000.0 * run_ds(cw); + tp_rg_ = now_ms() - r0; + last_n_disp_ = (int) ds.size(); + return rows_ ? finish_rows(Tn, tokens, out, t0) : finish(Tn, tokens, pos0, head, out, t0); +} + +bool Window::run_rows(const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err) { + if (S < 1 || S > MAXT || S > max_t_) { err = "metal window: batch window size out of range"; return false; } + if (active_ != 0) activate(0); + uint32_t* ri = rinfo_.as() + 8; + int64_t maxpos = 0; + for (int t = 0; t < S; ++t) { + if (rows[t] < 0 || rows[t] >= n_slots()) { err = "metal window: a batch row names a slot that does not exist"; return false; } + for (int u = 0; u < t; ++u) + if (rows[u] == rows[t]) { err = "metal window: two rows of one batch window are the same slot"; return false; } + const int set = rows[t] + 1; + rows_ss_[t] = sets_[(size_t) set]->ss; + rows_set_[t] = set; + ri[4 * t] = (uint32_t) set; + ri[4 * t + 1] = (uint32_t) pos[t]; + ri[4 * t + 2] = (uint32_t) sets_[(size_t) set]->hp; + ri[4 * t + 3] = 0; + maxpos = std::max(maxpos, pos[t]); + } + rows_ = true; + rows_n_ = S; + const bool ok = run(S, tokens, maxpos, true, out, err); + rows_ = false; + return ok; +} + +// the end of a window of batch slots: the picks, and each slot's commit (a one-token step committed its row: the other history +// parity is the committed one; the PLE history takes the row's normalized values; the session's token pair advances) +bool Window::finish_rows(int Tn, const int32_t* tokens, int32_t* out, double t0) { + ++win_n_[Tn]; + win_ms_[Tn] += now_ms() - t0; + win_gpu_[Tn] += last_gpu_ms_; + for (int t = 0; t < Tn; ++t) out[t] = tokout_.as()[t]; + for (int t = 0; t < Tn; ++t) { + SeqSet& s = *sets_[(size_t) rows_set_[t]]; + s.hp = 1 - s.hp; + if (ple_on_) { + const float* norm = normW_.as() + (size_t) t * strata::kernels::NG_HC_DIM; + for (int c = 0; c < strata::kernels::NG_HC_DIM; ++c) { + float* col = s.pleHist.as() + (size_t) c * strata::kernels::NG_HIST; + for (int r = 0; r + 1 < strata::kernels::NG_HIST; ++r) col[r] = col[r + 1]; + col[strata::kernels::NG_HIST - 1] = norm[c]; + } + } + s.ss->ple_prev[0] = s.ss->ple_prev[1]; + s.ss->ple_prev[1] = tokens[t]; + } + return true; +} + +double Window::run_ds(Cached& c) { + if (!stream_) return cx_->run_graph(c.ds); + return cx_->run_graph_seg(c.ds, c.cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this); +} + +namespace { +Window* g_stream_window = nullptr; +// What the process holds now, all of it (task_info TASK_VM_INFO): phys_footprint is its anonymous + wired + compressed memory +// (not the clean file pages it only maps), resident_size is the RSS (those pages too). Bytes; 0 when the call fails. +struct ProcMem { + double footprint = 0, rss = 0; +}; +ProcMem proc_mem() { + task_vm_info_data_t ti; + mach_msg_type_number_t cnt = TASK_VM_INFO_COUNT; + ProcMem r; + if (task_info(mach_task_self(), TASK_VM_INFO, (task_info_t) &ti, &cnt) == KERN_SUCCESS) { + r.footprint = (double) ti.phys_footprint; + r.rss = (double) ti.resident_size; + } + return r; +} +void stream_atexit() { + if (g_stream_window) g_stream_window->stream_report(stderr); +} +} // namespace + +// The arena, once the rest is allocated (the MTP layer loads after init): what the GPU working set and the RAM leave +// with STRATA_METAL_KEEP_FREE_GB kept for macOS (default: 7 GB, or 15% of the RAM when that is more - the M5 Max +// froze with an 89 GB arena and 7 GB left beside a 50 GB n-gram table read through the file cache; arena_budget.h); +// STRATA_METAL_ARENA_GB sets it. +bool Window::stream_arena(std::string& err) { + uint64_t phys = 0; + size_t len = sizeof phys; + sysctlbyname("hw.memsize", &phys, &len, nullptr, 0); + const double gb = 1e9; + if (!(std::getenv("STRATA_METAL_KEEP_HOST_WEIGHTS") && std::string(std::getenv("STRATA_METAL_KEEP_HOST_WEIGHTS")) == "1")) { + const uint64_t g = strata::gpu::purge_device_allocations(32u << 20); + std::fprintf(stderr, "strata metal stream: %.2f GB of the host code's device weight copies given back (unused by the " + "Metal window)\n", g / 1e9); + } + const double used = (double) cx_->allocated(); + // what macOS gives now without paging anything out (arena_avail: free, purgeable, the clean file cache; not other + // programs' inactive memory, freed only by compressing or swapping it) + auto avail_now = [&]() -> double { + vm_statistics64_data_t vs; + mach_msg_type_number_t cnt = HOST_VM_INFO64_COUNT; + if (host_statistics64(mach_host_self(), HOST_VM_INFO64, (host_info64_t) &vs, &cnt) == KERN_SUCCESS) { + const double pg = (double) vm_page_size; + return arena_avail(vs.free_count * pg, vs.purgeable_count * pg, vs.inactive_count * pg, vs.speculative_count * pg, + vs.external_page_count * pg); + } + return (double) phys - (double) cx_->allocated(); + }; + ArenaMem am; // arena_budget.h: the rule, testable + am.phys = (double) phys; + am.working_set = (double) cx_->recommended_working_set(); + am.allocated = used; + const ProcMem pm0 = proc_mem(); // everything the process holds, not only the Metal buffers (the old 67% rule missed it) + am.proc_footprint = pm0.footprint; + { // the file cache macOS keeps before it compresses anonymous memory (a 4-byte count of pages, read into a zeroed 8) + unsigned long long pages = 0; + size_t pl = sizeof pages; + if (sysctlbyname("vm.vm_page_filecache_min", &pages, &pl, nullptr, 0) == 0) am.filecache_min = (double) pages * (double) vm_page_size; + } + am.available = avail_now(); + if (main_ss_ && main_ss_->ple.table && main_ss_->ple.table->mode() == strata::kernels::PleIo::Mmap) + am.ple_mapped = (double) ss_->ple.table->rows() * 192.0; // a row is at most 192 bytes (Q8_0 + its scale) + if (const char* e = std::getenv("STRATA_METAL_KEEP_FREE_GB")) am.keep_gb = std::atof(e); + if (const char* e = std::getenv("STRATA_METAL_ARENA_GB")) am.arena_gb = std::atof(e); + if (const char* e = std::getenv("STRATA_METAL_LEAVE_SHARE")) am.leave_share = std::atof(e); // for measuring; the default is kLeaveShare + const double keep = arena_keep(am); + const double avail = am.available; + double budget = arena_budget(am); + if (budget < 2 * gb) { + std::fprintf(stderr, "strata metal stream: only %.1f GB left for the expert arena (the GPU working set %.1f GB, RAM %.1f GB, " + "%.1f GB allocated): expect a low hit rate\n", budget / gb, cx_->recommended_working_set() / gb, + phys / gb, used / gb); + budget = std::max(budget, 1.0 * gb); + } + std::string prof; + if (const char* e = std::getenv("STRATA_METAL_EXPERT_PROFILE")) prof = e; + else prof = strata_metal::resource(nullptr, "data", STRATA_METAL_DATA_DEFAULT) + "/expert-profile.bin"; + std::fprintf(stderr, "strata metal stream: %.2f GB allocated before the arena (the process holds %.2f GB, %.2f GB resident); " + "RAM %.1f GB (%.1f GB available, %.1f GB kept for macOS), GPU working set %.1f GB\n", used / gb, + pm0.footprint / gb, pm0.rss / gb, phys / gb, avail / gb, keep / gb, cx_->recommended_working_set() / gb); + { // the planned total: arena + the Metal buffers (or the process's own count when it is more) + the n-gram table's + // reserve, against the cap (arena_budget.h: footprint_cap) + const FootprintPlan fp = footprint_plan(am, budget); + std::fprintf(stderr, "strata metal stream: planned resident footprint %.1f GB = arena %.1f + %.1f held before it + %.1f " + "reserved for the n-gram table; cap %.1f GB = the lesser of the GPU working set %.1f GB and the RAM " + "%.1f GB less %.1f GB left to macOS and its file cache%s\n", + fp.total / gb, fp.arena / gb, fp.base / gb, fp.ple / gb, fp.cap / gb, cx_->recommended_working_set() / gb, + phys / gb, (am.phys - fp.cap) / gb, + am.arena_gb >= 0 && fp.over() ? " - STRATA_METAL_ARENA_GB is above the cap: kept as set (your choice), but " + "a Mac that runs short of memory can freeze (watchdog panics seen at 90+ GB " + "resident of 128)" : ""); + } + const char* epg = std::getenv("STRATA_METAL_PROMPT_GEMM"); + const bool pg_on = !(epg && std::string(epg) == "0"); + bool minimal = false; // the arena's least layout (the min slots) is not shrunk further: it goes ahead, filled in steps + for (int tries = 0;; ++tries) { + const double planned = budget; + std::function recheck = [&]() -> uint64_t { return (uint64_t) arena_shrink(am, planned, avail_now()); }; + if (st_->make_arena(*cx_, (uint64_t) budget, prof, pg_on ? 576 : 64, err, minimal ? std::function() : recheck, + [&](uint64_t step) { return arena_grow(am, avail_now(), (double) step); })) + break; + if (st_->shrink_to() == 0 || tries == 4) return false; + if ((double) st_->shrink_to() + 64.0 * (1 << 20) >= budget) minimal = true; // no smaller layout than this one + std::fprintf(stderr, "strata metal stream: less memory is free than when the arena was planned: %.1f -> %.1f GB\n", + budget / gb, st_->shrink_to() / gb); + budget = (double) st_->shrink_to(); + } + if (pg_on && !prompt_gemm_init(err)) return false; + { // what the process really holds now, against the plan and the cap (a number that disagrees is a bug in the count) + const ProcMem pm1 = proc_mem(); + const double cap = footprint_cap(am); + std::fprintf(stderr, "strata metal stream: resident footprint after the fill: the process holds %.1f GB (%.1f GB resident), " + "planned %.1f GB, cap %.1f GB%s\n", pm1.footprint / gb, pm1.rss / gb, + footprint_plan(am, (double) st_->arena_used_bytes()).total / gb, cap / gb, + pm1.footprint > cap ? " - OVER THE CAP" : ""); + } + g_stream_window = this; + std::atexit(stream_atexit); + return true; +} + +void Window::logpos(int Tn, const int32_t* tokens, int64_t pos0, bool head) { + static const char* path = std::getenv("STRATA_METAL_LOGPOS"); + if (path == nullptr || !head) return; + if (logpos_ == nullptr) logpos_ = std::fopen(path, "ab"); + if (logpos_ == nullptr) return; + if (lp_pend_ && Tn >= 1 && pos0 == lp_pos_ + 1) { + const int32_t tgt = tokens[0]; + std::fprintf(logpos_, "%lld\t%d\t%.6f\t%d\t0\t0\t0\t0", (long long) (lp_pos_ - lp_first_), tgt, + (double) lp_row_[(size_t) tgt] - lp_lse_, lp_arg_); + for (const auto& [v, id] : lp_top_) std::fprintf(logpos_, "\t%d:%.6f", id, (double) v - lp_lse_); + std::fputc('\n', logpos_); + std::fflush(logpos_); + } + lp_pend_ = false; + if (Tn != 1) return; + if (lp_first_ < 0) lp_first_ = pos0; + const float* lg = lm_.as(); + lp_row_.assign(lg, lg + n_vocab_); + float mx = lg[0]; + int32_t am = 0; + for (int64_t i = 1; i < n_vocab_; ++i) if (lg[i] > mx) { mx = lg[i]; am = (int32_t) i; } + double s = 0.0; + for (int64_t i = 0; i < n_vocab_; ++i) s += std::exp((double) (lg[i] - mx)); + lp_lse_ = (double) mx + std::log(s); + lp_arg_ = am; + lp_top_.clear(); + for (int64_t i = 0; i < n_vocab_; ++i) lp_top_.push_back({lg[i], (int32_t) i}); + std::partial_sort(lp_top_.begin(), lp_top_.begin() + 64, lp_top_.end(), + [](const auto& a, const auto& b) { return a.first > b.first || (a.first == b.first && a.second < b.second); }); + lp_top_.resize(64); + lp_pos_ = pos0; + lp_pend_ = true; +} + +void Window::kprof(const std::vector& ds, int Tn) { + std::map pname; + for (const auto& [k, v] : pcache_) pname[&v] = k.substr(0, k.find('#')); + std::map tname; + for (const auto& [k, v] : f_.tensors) { + std::string n = k; + if (n.rfind("blk.", 0) == 0) n = n.substr(n.find('.', 4) + 1); + tname[v.offset] = n; + } + std::map q8name; + for (const auto& [k, v] : q8off_) { + std::string n = k; + if (n.rfind("blk.", 0) == 0) n = n.substr(n.find('.', 4) + 1); + q8name[v] = n; + } + std::map> by; + for (const auto& d : ds) { + std::string c = pname.count(d.p) ? pname[d.p] : "?"; + if (!d.args.empty() && d.args[0].buf == &W_ && tname.count(d.args[0].off) && c.rfind("moe_", 0) != 0) + c += " " + tname[d.args[0].off]; + else if (!d.args.empty() && d.args[0].buf == &Q8_) c += " " + (q8name.count(d.args[0].off) ? q8name[d.args[0].off] : std::string("Q8")); + mtl::Dispatch x = d; + x.barrier = true; + by[c].push_back(x); + } + auto med = [&](const std::vector& v) { + std::vector t; + for (int i = 0; i < 7; ++i) t.push_back(1000.0 * cx_->run_graph(v)); + std::sort(t.begin(), t.end()); + return t[3]; + }; + const double all = med(ds); + std::vector> rows; + double sum = 0; + for (const auto& [c, v] : by) { + const double m = med(v); + sum += m; + rows.push_back({m, c + " x" + std::to_string(v.size())}); + } + std::sort(rows.rbegin(), rows.rend()); + std::fprintf(stderr, "strata metal kprof: window T%d, %zu dispatches, GPU %.3f ms; classes alone (serialized) sum %.3f ms:\n", + Tn, ds.size(), all, sum); + for (const auto& [m, c] : rows) std::fprintf(stderr, "strata metal kprof: %8.3f ms %5.1f%% %s\n", m, 100.0 * m / sum, c.c_str()); +} + +Window::~Window() { + if (logpos_) std::fclose(logpos_); + report_profile(); + if (pb_.lat_gu16.handle) cx_->free(pb_.lat_gu16); +} + +void Window::report_profile() const { + if (std::getenv("STRATA_METAL_PROFILE") == nullptr) return; + auto line = [](const char* what, const Prof& x) { + if (x.n == 0) return; + std::fprintf(stderr, "strata metal profile: %-10s %6lld x %.3f ms: host before %.3f, run_graph %.3f (GPU %.3f, so " + "encode/submit/wait %.3f), patch+finish %.3f\n", what, (long long) x.n, x.all / x.n, x.pre / x.n, + x.rg / x.n, x.gpu / x.n, (x.rg - x.gpu) / x.n, x.fin / x.n); + }; + for (const auto& [k, x] : prof_) { + const std::string w = "window T" + std::to_string(k / 2) + (k % 2 ? "" : " nohead"); + line(w.c_str(), x); + } + line("commit", cprof_); + line("draft", dprof_); +} + +bool Window::finish(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, double t0) { + ++win_n_[Tn]; + win_ms_[Tn] += now_ms() - t0; + win_gpu_[Tn] += last_gpu_ms_; + static const bool prof = std::getenv("STRATA_METAL_PROFILE") != nullptr; + if (prof) { + Prof& x = prof_[Tn * 2 + (head ? 1 : 0)]; + const double t1 = now_ms(); + ++x.n; x.pre += tp_pre_ - t0; x.rg += tp_rg_; x.gpu += last_gpu_ms_; x.fin += t1 - tp_pre_ - tp_rg_; + x.all += t1 - t0; + } + logpos(Tn, tokens, pos0, head); + if (Tn == 1 || partial_) hp_ = 1 - hp_; // the window committed its first row (or all of them) in the step + last_head_ = head; + for (int t = 0; t < Tn; ++t) out[t] = head ? tokout_.as()[t] : -1; + last_T_ = Tn; + for (int t = 0; t < Tn; ++t) last_tokens_[t] = tokens[t]; + static const bool trace = std::getenv("STRATA_METAL_TRACE") != nullptr; + if (trace && head) { + // the pick and its margin over the runner-up (a near tie says a backend difference may flip it) + const float* lg = lm_.as(); + int i1 = -1, i2 = -1; + for (int i = 0; i < n_vocab_; ++i) { + if (i1 < 0 || lg[i] > lg[i1]) { i2 = i1; i1 = i; } + else if (i2 < 0 || lg[i] > lg[i2]) i2 = i; + } + std::fprintf(stderr, "strata metal: window pos %lld token %d -> %d (top %d %.4f, 2nd %d %.4f, margin %.4f) %.2f ms (GPU %.2f ms, %d dispatches)\n", + (long long) pos0, tokens[0], out[0], i1, lg[i1], i2, lg[i2], lg[i1] - lg[i2], now_ms() - t0, last_gpu_ms_, last_n_disp_); + } + return true; +} + +void Window::fold(std::vector& ds, int n_keep, int hp) { + const size_t off = partial_ ? 1 : 0; // rows 1.. (row 0 was committed in the window's step) + for (int l = 0; l < 48; ++l) { + if (l % 4 == 3) continue; + const int gi = l - l / 4; + const std::string b = "blk." + std::to_string(l) + "."; + const size_t o_qkv = ((size_t) gi * MAXT + off) * 10240 * 4, o_ab = ((size_t) gi * MAXT + off) * 48 * 4, + o_z = ((size_t) gi * MAXT + off) * 6144 * 4; + GdnArgs a{16, 48, 10240, 1u, eps_, 0}; + mtl::Dispatch d; + d.p = P(lgdn_, "gdn_step", mtl::Constants().i(2, n_keep), "gdn_step#" + std::to_string(n_keep)); + d.args = {mtl::Arg::B(gst_[(size_t) gi]), mtl::Arg::B(ghist_[hp][(size_t) gi]), mtl::Arg::B(qkvL_, o_qkv), + mtl::Arg::B(alL_, o_ab), mtl::Arg::B(beL_, o_ab), mtl::Arg::B(zL_, o_z), + mtl::Arg::B(W_, T(b + "ssm_conv1d.weight").offset), mtl::Arg::B(W_, T(b + "ssm_dt.bias").offset), + mtl::Arg::B(W_, T(b + "ssm_a").offset), mtl::Arg::B(W_, T(b + "ssm_norm.weight").offset), + mtl::Arg::B(ydump_), mtl::Arg::V(a), mtl::Arg::B(ghist_[1 - hp][(size_t) gi])}; + d.grid[0] = 48; + d.tg[0] = 512; + nb_ = true; // the layers' folds are independent of each other + push(ds, d); + } + nb_ = false; // what follows reads the folded state: a barrier +} + +void Window::swap_state(SeqSet& o) { + auto sw = [](std::vector& a, std::vector& b) { + for (size_t i = 0; i < a.size() && i < b.size(); ++i) std::swap(a[i], b[i]); + }; + sw(gst_, o.gst); sw(ghist_[0], o.gh0); sw(ghist_[1], o.gh1); sw(kq_, o.kq); sw(vq_, o.vq); sw(ksc_, o.ksc); sw(vsc_, o.vsc); + sw(tail_, o.tail); sw(dead_, o.dead); sw(pooled_, o.pooled); + std::swap(pleHist_, o.pleHist); + ple_hist_.swap(o.ple_hist); + last_norm_.swap(o.last_norm); + std::swap(hp_, o.hp); std::swap(pend_fold_, o.pend_fold); std::swap(last_T_, o.last_T); std::swap(last_head_, o.last_head); + for (int t = 0; t < MAXT; ++t) std::swap(last_tokens_[t], o.last_tokens[t]); + std::swap(ss_, o.ss); +} + +// the buffers of one sequence's running state, shaped like the live ones (real: allocated and zero; else empty placeholders) +void Window::alloc_set(SeqSet& o, bool real, int idx) { + const size_t MAXC = (size_t) max_cells_, MB = MAXC / 4 + 2; + if (real && bigs_ok_) { // set idx's part of the big buffers + for (size_t i = 0; i < gst_.size(); ++i) { + o.gst.push_back(view_of(gstBig_[i], (size_t) idx * st_gst_, 4 * 128 * 48 * 128)); + o.gh0.push_back(view_of(hbBig_[i], ((size_t) idx * 2 + 0) * st_hist_, 4 * 10240 * 3)); + o.gh1.push_back(view_of(hbBig_[i], ((size_t) idx * 2 + 1) * st_hist_, 4 * 10240 * 3)); + } + for (size_t j = 0; j < kq_.size(); ++j) { + o.kq.push_back(view_of(kqBig_[j], (size_t) idx * st_kq_, MAXC * 2 * 256)); + o.vq.push_back(view_of(vqBig_[j], (size_t) idx * st_kq_, MAXC * 2 * 256)); + o.ksc.push_back(view_of(kscBig_[j], (size_t) idx * st_sc_, 2 * MAXC * 2 * 4)); + o.vsc.push_back(view_of(vscBig_[j], (size_t) idx * st_sc_, 2 * MAXC * 2 * 4)); + o.tail.push_back(view_of(tailBig_[j], (size_t) idx * st_tail_, 4 * 4 * 128)); + o.dead.push_back(view_of(deadBig_[j], (size_t) idx * st_dead_, 4 * 128)); + o.pooled.push_back(view_of(pooledBig_[j], (size_t) idx * st_pooled_, 4 * MB * 128)); + } + o.pleHist = view_of(pleBig_, (size_t) idx * st_ple_, 4 * 9 * 10240); + o.ple_hist.assign(ple_hist_.size(), 0.0f); + return; + } + auto A = [&](size_t bytes) { return real ? cx_->alloc(bytes) : mtl::Buffer(); }; + for (size_t i = 0; i < gst_.size(); ++i) { + o.gst.push_back(A(4 * 128 * 48 * 128)); + o.gh0.push_back(A(4 * 10240 * 3)); + o.gh1.push_back(A(4 * 10240 * 3)); + } + for (size_t j = 0; j < kq_.size(); ++j) { + o.kq.push_back(A(MAXC * 2 * 256)); + o.vq.push_back(A(MAXC * 2 * 256)); + o.ksc.push_back(A(2 * MAXC * 2 * 4)); + o.vsc.push_back(A(2 * MAXC * 2 * 4)); + o.tail.push_back(A(4 * 4 * 128)); + o.dead.push_back(A(4 * 128)); + o.pooled.push_back(A(4 * MB * 128)); + } + o.pleHist = A(4 * 9 * 10240); + if (real) o.ple_hist.assign(ple_hist_.size(), 0.0f); +} + +bool Window::slots_init(const std::vector& slots, std::string& err) { + if (!sets_.empty()) return true; // (asked once) + if (cx_ == nullptr) { err = "metal window: slots before the window is up"; return false; } + std::vector> v; + v.push_back(std::make_unique()); + alloc_set(*v[0], false); + v[0]->ss = ss_; + set_of_ss_[ss_] = 0; + uint64_t bytes = 0; + for (size_t b = 0; b < slots.size(); ++b) { + if ((int) b >= g_slot_cap) bigs_ok_ = false; // more slots than the window was sized for: serial windows only + auto x = std::make_unique(); + alloc_set(*x, true, (int) b + 1); + x->ss = slots[b]; + for (auto* vec : {&x->gst, &x->gh0, &x->gh1, &x->kq, &x->vq, &x->ksc, &x->vsc, &x->tail, &x->dead, &x->pooled}) + for (const auto& bf : *vec) bytes += bf.bytes; + set_of_ss_[slots[b]] = (int) b + 1; + v.push_back(std::move(x)); + } + sets_ = std::move(v); + active_ = 0; + if (bigs_ok_) { // the strides the row kernels need (elements of each array's type) + uint32_t* h = rinfo_.as(); + h[0] = (uint32_t) st_kq_; h[1] = (uint32_t) (st_sc_ / 2); h[2] = (uint32_t) (st_tail_ / 4); h[3] = (uint32_t) (st_dead_ / 4); + h[4] = (uint32_t) (st_pooled_ / 4); h[5] = (uint32_t) (st_gst_ / 4); h[6] = (uint32_t) (st_hist_ / 4); h[7] = (uint32_t) (st_ple_ / 4); + } + std::fprintf(stderr, "strata metal: %zu batch slots, each with its own running state and K/V (%.2f GB each, %.2f GB in all)\n", + slots.size(), slots.empty() ? 0.0 : (double) bytes / (double) slots.size() / 1e9, (double) bytes / 1e9); + return true; +} + +// Makes `set` the live state. The live members and the stashes exchange contents (two swaps through the active set's own stash). +void Window::activate(int set) { + if (set == active_ || set < 0 || set >= (int) sets_.size()) return; + swap_state(*sets_[(size_t) active_]); // the stash of the active one now holds its state, the live members the spare + swap_state(*sets_[(size_t) set]); // ... and now the wanted one + active_ = set; +} + +bool Window::run_slots(const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err) { + if (sets_.empty()) { err = "metal window: no batch slots"; return false; } + static const bool one_window = !(std::getenv("STRATA_METAL_BATCH_ROWS") && std::string(std::getenv("STRATA_METAL_BATCH_ROWS")) == "0"); + if (one_window && bigs_ok_) { + activate(0); + return run_rows(rows, S, tokens, pos, out, err); + } + bool ok = true; + for (int t = 0; t < S && ok; ++t) { + if (rows[t] < 0 || rows[t] >= n_slots()) { err = "metal window: a batch row names a slot that does not exist"; ok = false; break; } + activate(rows[t] + 1); + // a one-token window commits its row in its own step; commit(1) adds the PLE history and the session's token pair, + // which read buffers the next window rewrites + ok = run(1, tokens + t, pos[t], true, out + t, err) && commit(1, err); + } + activate(0); + return ok; +} + +bool Window::slot_copy(bool into_slot, int b, int64_t upto, bool kv_only, std::string& err) { + if (b < 0 || b >= n_slots()) { err = "metal window: no such batch slot"; return false; } + activate(0); + SeqSet& o = *sets_[(size_t) b + 1]; + if (pend_fold_ > 0 || o.pend_fold > 0) { err = "metal window: a slot copy with a fold pending"; return false; } + const size_t cells = (size_t) std::max(0, std::min(upto, max_cells_)); + auto cp = [&](const mtl::Buffer& live, const mtl::Buffer& stash, size_t n) { + n = std::min({n, live.bytes, stash.bytes}); + if (into_slot) std::memcpy(stash.ptr, live.ptr, n); else std::memcpy(live.ptr, stash.ptr, n); + }; + const int hs = into_slot ? hp_ : o.hp; // the committed history parity of the source + if (!kv_only) { + if (into_slot) o.hp = hp_; else hp_ = o.hp; + for (size_t i = 0; i < gst_.size(); ++i) { + cp(gst_[i], o.gst[i], gst_[i].bytes); + cp(ghist_[hs][i], hs ? o.gh1[i] : o.gh0[i], ghist_[hs][i].bytes); + } + } + for (size_t j = 0; j < kq_.size(); ++j) { + cp(kq_[j], o.kq[j], cells * 512); + cp(vq_[j], o.vq[j], cells * 512); + cp(ksc_[j], o.ksc[j], cells * 16); + cp(vsc_[j], o.vsc[j], cells * 16); + cp(pooled_[j], o.pooled[j], ((cells + 3) / 4 + 2) * 512); + if (kv_only) continue; + cp(tail_[j], o.tail[j], tail_[j].bytes); + cp(dead_[j], o.dead[j], dead_[j].bytes); + } + if (kv_only) return true; + cp(pleHist_, o.pleHist, pleHist_.bytes); + if (into_slot) { o.ple_hist = ple_hist_; o.ss->ple_prev[0] = ss_->ple_prev[0]; o.ss->ple_prev[1] = ss_->ple_prev[1]; } + else { ple_hist_ = o.ple_hist; ss_->ple_prev[0] = o.ss->ple_prev[0]; ss_->ple_prev[1] = o.ss->ple_prev[1]; } + return true; +} + +bool Window::checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err) { + (void) err; + if (pend_fold_ > 0) { // the last commit's fold first: the checkpoint is the state after it + std::vector ds; + nb_ = false; + fold(ds, pend_fold_, 1 - hp_); + cx_->run_graph(ds); + pend_fold_ = 0; + } + const size_t gs = 4 * 128 * 48 * 128, hs = 4 * 10240 * 3, ts = 4 * 4 * 128, dsz = 4 * 128; + c.gdn.resize(gst_.size() * (gs + hs)); + for (size_t i = 0; i < gst_.size(); ++i) { + std::memcpy(c.gdn.data() + i * (gs + hs), gst_[i].ptr, gs); + std::memcpy(c.gdn.data() + i * (gs + hs) + gs, ghist_[hp_][i].ptr, hs); + } + c.ple.assign((const uint8_t*) pleHist_.ptr, (const uint8_t*) pleHist_.ptr + 4 * 9 * 10240); + c.tails.resize(tail_.size() * ts); + c.dead.resize(dead_.size() * dsz); + for (size_t j = 0; j < tail_.size(); ++j) { + std::memcpy(c.tails.data() + j * ts, tail_[j].ptr, ts); + std::memcpy(c.dead.data() + j * dsz, dead_[j].ptr, dsz); + } + c.block_pos.clear(); + return true; +} + +void Window::reset() { + pend_fold_ = 0; + last_T_ = 0; + auto zero = [](const mtl::Buffer& b, size_t n) { std::memset(b.ptr, 0, n); }; + for (size_t i = 0; i < gst_.size(); ++i) { + zero(gst_[i], 4 * 128 * 48 * 128); + zero(ghist_[0][i], 4 * 10240 * 3); + zero(ghist_[1][i], 4 * 10240 * 3); + } + zero(pleHist_, 4 * 9 * 10240); + for (size_t j = 0; j < tail_.size(); ++j) { + zero(tail_[j], 4 * 4 * 128); + zero(dead_[j], 4 * 128); + } + for (size_t j = 0; j < mtail_.size(); ++j) { + zero(mtail_[j], 4 * 4 * 128); + zero(mdead_[j], 4 * 128); + } +} + +bool Window::checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err) { + const size_t gs = 4 * 128 * 48 * 128, hs = 4 * 10240 * 3, ts = 4 * 4 * 128, dsz = 4 * 128; + if (c.gdn.size() != gst_.size() * (gs + hs) || c.ple.size() != (size_t) 4 * 9 * 10240 || + c.tails.size() != tail_.size() * ts || c.dead.size() != dead_.size() * dsz) { + err = "metal window: the checkpoint is not this engine's (sizes differ)"; + return false; + } + pend_fold_ = 0; // the restored state replaces whatever was pending + last_T_ = 0; + for (size_t i = 0; i < gst_.size(); ++i) { + std::memcpy(gst_[i].ptr, c.gdn.data() + i * (gs + hs), gs); + std::memcpy(ghist_[hp_][i].ptr, c.gdn.data() + i * (gs + hs) + gs, hs); + } + std::memcpy(pleHist_.ptr, c.ple.data(), c.ple.size()); + for (size_t j = 0; j < tail_.size(); ++j) { + std::memcpy(tail_[j].ptr, c.tails.data() + j * ts, ts); + std::memcpy(dead_[j].ptr, c.dead.data() + j * dsz, dsz); + } + return true; +} + +bool Window::commit(int n_keep, std::string& err) { + if (n_keep < 1 || n_keep > last_T_) { err = "metal window: commit count out of range"; return false; } + using namespace strata::kernels; + if (last_T_ > 1) { + // the GDN layers: the kept prefix folded into the state (gdn_step over n_keep tokens from the window's own + // inputs, with commit; bitwise the window's trajectory). K/V and indexer cells past the prefix are rewritten + // when those positions run again, before anything reads them. + static const bool lazy = !(std::getenv("STRATA_METAL_LAZY_FOLD") && std::string(std::getenv("STRATA_METAL_LAZY_FOLD")) == "0"); + int nf = n_keep; + if (partial_) { + if (!last_head_) { // the prompt's window committed every row in its step + if (n_keep != last_T_) { err = "metal window: a prompt window commits all its rows"; return false; } + nf = 0; + } else { + nf = n_keep - 1; // row 0 is in + } + } + if (nf > 0) { + if (lazy) { + pend_fold_ = nf; // folded at the start of the next window (from the history parity before the turn) + } else { + std::vector ds; + nb_ = false; + fold(ds, nf, hp_); + cx_->run_graph(ds); + } + hp_ = 1 - hp_; + } + } + strata::core::SessionState& ss = *ss_; + for (int t = 0; t < n_keep; ++t) { + if (ple_on_) { + const float* norm = normW_.as() + (size_t) t * NG_HC_DIM; + for (int c = 0; c < NG_HC_DIM; ++c) { + float* col = pleHist_.as() + (size_t) c * NG_HIST; + for (int r = 0; r + 1 < NG_HIST; ++r) col[r] = col[r + 1]; + col[NG_HIST - 1] = norm[c]; + } + } + ss.ple_prev[0] = ss.ple_prev[1]; + ss.ple_prev[1] = last_tokens_[t]; + } + return true; +} + + +// ================================================================================================ the MTP draft layer +bool Window::mtp_load(const std::string& gguf_path, int max_t, std::string& err) { + (void) max_t; + if (!fm_.open(gguf_path, err)) return false; + madvise(fm_.map, fm_.map_bytes, MADV_WILLNEED); + Wm_ = cx_->wrap(fm_.map, fm_.map_bytes); + if (Wm_.handle == nullptr) { err = "metal mtp: cannot wrap " + gguf_path; return false; } + // vLLM's GemmaRMSNorm scales by (1 + w): the F32 norms folded once (tools/mtp_rt.py does the same for CUDA) + size_t tot = 0; + for (const auto& kv : fm_.tensors) + if (kv.second.type == 0) { mnorm_off_[kv.first] = tot; tot += (kv.second.bytes + 255) / 256 * 256; } + mnorm_ = cx_->alloc(tot + 256); + for (const auto& kv : mnorm_off_) { + const gguf::Tensor& x = fm_.tensors.at(kv.first); + const float* src = (const float*) ((const uint8_t*) fm_.map + x.offset); + float* dst = (float*) (mnorm_.as() + kv.second); + for (uint64_t i = 0; i < x.bytes / 4; ++i) dst[i] = src[i] + 1.0f; + } + lmtp_ = cx_->compile({strata_metal::resource("STRATA_METAL_ENGINE_KERNELS", "engine-msl", STRATA_METAL_ENGINE_KERNELS_DEFAULT) + "/mtp.metal"}, "#define NO_GGML_COMMON 1\n"); + { + const char* e = std::getenv("STRATA_METAL_MTP_Q8"); + if (!(e && std::string(e) == "0")) { + size_t tot = 0; + std::vector names; + for (const auto& kv : fm_.tensors) { + const gguf::Tensor& x = kv.second; + uint64_t ne = 1; + for (uint64_t d : x.shape) ne *= d; + if (x.type != ref::BF16 || x.shape.size() < 2 || x.shape[0] % 32 != 0) continue; + gguf::Tensor q = x; + q.type = ref::Q8_0; + q.offset = tot; + q.bytes = ne / 32 * 34; + mq8t_[x.name] = q; + names.push_back(x.name); + tot += (q.bytes + 255) / 256 * 256; + } + mq8_ = cx_->alloc(tot + 4096); + std::vector th; + const unsigned nth = std::max(1u, std::min(16u, std::thread::hardware_concurrency())); + for (unsigned k = 0; k < nth; ++k) + th.emplace_back([&, k] { + std::vector tmp; + for (size_t i = k; i < names.size(); i += nth) { + const gguf::Tensor& x = fm_.tensors.at(names[i]); + const uint64_t ne = x.bytes / 2; + tmp.resize(ne); + ref::dequant(ref::BF16, (const uint8_t*) fm_.map + x.offset, tmp.data(), ne); + ref::quantize_q8_0(tmp.data(), ne, mq8_.as() + mq8t_[names[i]].offset); + } + }); + for (auto& x : th) x.join(); + } + } + auto A = [&](size_t bytes) { return cx_->alloc(bytes); }; + const int MAXC = (int) max_cells_, MB = MAXC / 4 + 2; + mR_ = A(4 * MAXT * HC * N); memb_ = A(4 * MAXT * N); men_ = A(4 * MAXT * N); me2_ = A(4 * MAXT * N); + mhn_ = A(4 * MAXT * HC * N); mh2_ = A(4 * HC * MAXT * N); midx_ = A(4 * MAXT * 640); mbo_ = A(4 * MAXT * N); + minj_[0] = A(4 * MAXT * HC); minj_[1] = A(4 * MAXT * HC); mqids_ = A(4 * MAXT * 2051); mlm_ = A(4 * (size_t) n_vocab_); + auto act = [&](int n) { Act x; x.n = n; x.b = A(4 * ((size_t) 4 * MAXT * n + (size_t) MAXT * n / 16)); return x; }; + XEn_ = act(N); XHn_ = act(N); + mkq_.push_back(A((size_t) MAXC * 2 * 256)); mvq_.push_back(A((size_t) MAXC * 2 * 256)); + mksc_.push_back(A(2 * (size_t) MAXC * 2 * 4)); mvsc_.push_back(A(2 * (size_t) MAXC * 2 * 4)); + mtail_.push_back(A(4 * 4 * 128)); mdead_.push_back(A(4 * 128)); mpooled_.push_back(A(4 * (size_t) MB * 128)); + { + std::string vp; + // the subset setup put beside the draft layer (--draft-vocab: /rt/draft_vocab.bin, as CUDA reads it), + // else the repository's default + const std::string beside = gguf_path.substr(0, gguf_path.find_last_of('/') + 1) + "rt/draft_vocab.bin"; + if (const char* e = std::getenv("STRATA_METAL_DRAFT_VOCAB")) vp = e; + else if (FILE* t = std::fopen(beside.c_str(), "rb")) { std::fclose(t); vp = beside; } + else vp = strata_metal::resource(nullptr, "data", STRATA_METAL_DATA_DEFAULT) + "/draft_vocab.bin"; + FILE* f = vp == "full" ? nullptr : std::fopen(vp.c_str(), "rb"); + if (f) { + int32_t id; + while (std::fread(&id, 4, 1, f) == 1) if (id >= 0 && id < n_vocab_) dvocab_.push_back(id); + std::fclose(f); + } + if (!dvocab_.empty()) { + const gguf::Tensor& h = T("output.weight"); + const size_t rb = ref::row_bytes(h.type, N); + mhead_ = A(dvocab_.size() * rb + 4096); + for (size_t i = 0; i < dvocab_.size(); ++i) + std::memcpy(mhead_.as() + i * rb, (const uint8_t*) W_.ptr + h.offset + (size_t) dvocab_[i] * rb, rb); + mhead_t_ = h; + mhead_t_.offset = 0; + mhead_t_.shape = {(uint64_t) N, (uint64_t) dvocab_.size()}; + mhead_t_.bytes = dvocab_.size() * rb; + } + } + mtp_on_ = true; + { + const int ty = T("token_embd.weight").type; + const int ok_types[] = {30, 8, 2, 6, 20, 23, 42, 21, 11, 12, 13, 14, 18}; + const char* e = std::getenv("STRATA_METAL_MTP_CHAIN"); + chain_ = !(e && std::string(e) == "0") && std::find(std::begin(ok_types), std::end(ok_types), ty) != std::end(ok_types); + if (chain_) { + lchain_ = cx_->compile({"quant.metal", "units.metal", strata_metal::resource("STRATA_METAL_ENGINE_KERNELS", "engine-msl", STRATA_METAL_ENGINE_KERNELS_DEFAULT) + "/mtp_chain.metal"}); + dout_ = A(4 * MAXT); dprob_ = A(4 * MAXT); posv_ = A(4 * N); + for (int i = 0; i < N; ++i) posv_.as()[i] = (float) i; + dmap_ = A(4 * std::max(1, dvocab_.size())); + if (!dvocab_.empty()) std::memcpy(dmap_.as(), dvocab_.data(), 4 * dvocab_.size()); + // parity: a few tokens' embeddings, GPU vs the host's dequant (the per-step path's) + const gguf::Tensor& te = T("token_embd.weight"); + const size_t rb = ref::row_bytes(te.type, N); + const int32_t probe[MAXT] = {0, 1, 271, 760, 6511, 11751, 100000, (int32_t) n_vocab_ - 1}; + double worst = 0.0; + std::vector ref_row(N); + for (int k = 0; k < MAXT; ++k) { + dout_.as()[0] = probe[k]; + std::vector ds; + nb_ = false; + mtl::Dispatch d; + d.p = P(lchain_, "mtp_embed_" + std::to_string(ty), {}, "mtp_embed_" + std::to_string(ty)); + struct EmbArgs { uint32_t n, row_bytes, j; } ea{(uint32_t) N, (uint32_t) rb, 0}; + d.args = {mtl::Arg::B(W_, te.offset), mtl::Arg::B(dout_), mtl::Arg::B(posv_), mtl::Arg::B(memb_), mtl::Arg::V(ea)}; + d.tg[0] = 256; + push(ds, d); + cx_->run_graph(ds); + ref::dequant(te.type, (const uint8_t*) W_.ptr + te.offset + (size_t) probe[k] * rb, ref_row.data(), N); + for (int i = 0; i < N; ++i) worst = std::max(worst, (double) std::fabs(ref_row[i] - memb_.as()[i])); + } + std::fprintf(stderr, "strata metal: MTP draft round in one command buffer (embedding type %d on the GPU: max " + "|GPU - host| %.3g over %d rows)\n", ty, worst, MAXT); + if (worst > 1e-6) { + std::fprintf(stderr, "strata metal: the GPU embedding differs from the host's: a command buffer per step\n"); + chain_ = false; + } + } + } + std::fprintf(stderr, "strata metal: MTP draft layer from %s (%.2f GB, %zu norms folded), draft head over %s tokens\n", + gguf_path.c_str(), fm_.file_bytes / 1e9, mnorm_off_.size(), + dvocab_.empty() ? "all" : std::to_string(dvocab_.size()).c_str()); + return true; +} + +void Window::gr_read_m(std::vector& ds, const std::string& pre, mtl::Buffer& Rb, const mtl::Buffer& bo, + const mtl::Buffer& inj_in, const mtl::Buffer& inj_out, bool apply, const Act& out, bool final_, + int ntok) { + const gguf::Tensor& tdn = TQ(pre + ".input_mix_weight_down.weight"); + const gguf::Tensor& tup = TQ(pre + ".input_mix_weight_up.weight"); + const gguf::Tensor& tin = final_ ? tdn : TQ(pre + ".block_inject_weight.weight"); + const bool q8 = tdn.type == ref::Q8_0 && tup.type == ref::Q8_0 && tin.type == ref::Q8_0; + const mtl::Buffer& WB = q8 ? mq8_ : Wm_; + GrArgs a{N, HC, LR, (uint32_t) ntok, apply ? 1u : 0u, eps_}; + mtl::Dispatch d1; + d1.p = P(lib_, "gr_norm", {}, "gr_norm"); + d1.args = {mtl::Arg::B(Rb), mtl::Arg::B(bo), mtl::Arg::B(inj_in), mtl::Arg::B(mnorm_, NM(pre + ".hc_norm.weight")), + mtl::Arg::B(xn_), mtl::Arg::V(a)}; + d1.grid[0] = HC; d1.grid[1] = ntok; d1.tg[0] = 256; + push(ds, d1); + mtl::Dispatch d2; + const char* dn = q8 ? "gr_down_q8_8_8" : "gr_down_8_8"; + d2.p = P(lib_, dn, mtl::Constants().i(1, ntok), std::string(dn) + "#" + std::to_string(ntok)); + d2.args = {mtl::Arg::B(WB, tdn.offset), mtl::Arg::B(WB, tin.offset), mtl::Arg::B(xn_), mtl::Arg::B(lo_), + mtl::Arg::B(inj_out), mtl::Arg::V(a)}; + d2.grid[0] = LR + HC; d2.tg[0] = 256; + push(ds, d2); + mtl::Dispatch d3; + const char* up = q8 ? "gr_up_q8" : "gr_up"; + d3.p = P(lib_, up, mtl::Constants().i(1, ntok), std::string(up) + "#" + std::to_string(ntok)); + GrArgs au = a; + au.n_tok = MAXT; + d3.args = {mtl::Arg::B(WB, tup.offset), mtl::Arg::B(lo_), mtl::Arg::B(xn_), mtl::Arg::B(out.b), mtl::Arg::V(au)}; + d3.grid[0] = N / 16; d3.tg[0] = 256; + push(ds, d3); +} + +void Window::mtp_encode(std::vector& ds, const mtl::Buffer& srcR, size_t src_off, int n, const int32_t* tokens, + int64_t cell0, bool full, int erow, int slot) { + struct RmsArgs { uint32_t n; float eps; }; + struct XpArgs { uint32_t n, n_tok, plane_tok, stride; }; + struct CombArgs { uint32_t n, hc, n_tok, h_stride; }; + // ---- the embedding rows (host tokens: dequantized on the CPU into the shared buffer) + const gguf::Tensor& te = T("token_embd.weight"); + const size_t rb = ref::row_bytes(te.type, N); + if (tokens) + for (int i = 0; i < n; ++i) + ref::dequant(te.type, (const uint8_t*) W_.ptr + te.offset + (size_t) tokens[i] * rb, + (membp_ ? *membp_ : memb_).as() + (size_t) (erow + i) * N, N); + const std::string L = "mtp.layers.0."; + auto rms = [&](const mtl::Buffer& x, size_t xo, const std::string& w, mtl::Buffer& y, int width, int rows) { + mtl::Dispatch d; + d.p = P(lmtp_, "mtp_rms", {}, "mtp_rms"); + RmsArgs a{(uint32_t) width, eps_}; + d.args = {mtl::Arg::B(x, xo), mtl::Arg::B(mnorm_, NM(w)), mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = rows; d.tg[0] = 256; + push(ds, d); + }; + auto xp = [&](const mtl::Buffer& x, size_t xo, const Act& out, int width, int rows, int stride) { + mtl::Dispatch d; + d.p = P(lmtp_, "mtp_xprep", {}, "mtp_xprep"); + XpArgs a{(uint32_t) width, (uint32_t) rows, MAXT, (uint32_t) stride}; + d.args = {mtl::Arg::B(x, xo), mtl::Arg::B(out.b), mtl::Arg::V(a)}; + d.grid[0] = (width / 16 * rows + 255) / 256; d.tg[0] = 256; + push(ds, d); + }; + auto gm = [&](const std::string& name, const Act& x, const mtl::Buffer& y, size_t y_stride, int rows, size_t y_off = 0) { + // a GEMV on the MTP model's tensor (the window's gemv() reads the main model's buffer) + const gguf::Tensor& tt = TQ(name); + const int n_in = (int) tt.shape[0], n_out = tt.shape.size() > 1 ? (int) tt.shape[1] : 1, ty = tt.type; + int Lq = 16, Rr = 1, S = 4, K = 1; + if (n_in >= 6144) { Lq = 32; Rr = 1; S = 8; K = 2; } + else if (n_in == 640) { Lq = 8; Rr = 1; S = 4; K = 1; } + else if (n_out <= 640 || ty == ref::Q2_0 || ty == ref::BF16) { Lq = 32; Rr = 1; S = 4; K = 1; } + const std::string fn = "gemv_" + std::to_string(ty) + "_" + std::to_string(Lq) + "_" + std::to_string(Rr) + "_" + + std::to_string(S) + "_" + std::to_string(K); + GemvArgs a{(uint32_t) n_in, (uint32_t) n_out, (uint32_t) rows, (uint32_t) ref::row_bytes(ty, n_in), (uint32_t) n_in, + (uint32_t) y_stride}; + mtl::Dispatch d; + d.p = P(lib_, fn, mtl::Constants().i(0, rows), fn + "#" + std::to_string(rows)); + d.args = {mtl::Arg::B(MB(name), tt.offset), mtl::Arg::B(x.b, x.plane_bytes() * xkind(ty)), + mtl::Arg::B(x.b, x.plane_bytes() * 4), mtl::Arg::B(y, y_off), mtl::Arg::V(a)}; + const int rpt = (S / K) * (32 / Lq) * Rr; + d.grid[0] = (n_out + rpt - 1) / rpt; d.tg[0] = S * 32; + push(ds, d); + }; + // ---- e = fc_embedding(rms(embed(tok)) * (1 + w)), h = fc_hidden per stream (rms_10240(R) * (1 + w)), R = h + e + rms(membp_ ? *membp_ : memb_, (size_t) erow * N * 4, "mtp.pre_fc_norm_embedding.weight", men_, N, n); + xp(men_, 0, XEn_, N, n, N); + gm("mtp.fc_embedding.weight", XEn_, me2_, N, n); + rms(srcR, src_off, "mtp.pre_fc_norm_hidden.weight", mhn_, HC * N, n); + for (int c = 0; c < HC; ++c) { + xp(mhn_, (size_t) c * N * 4, XHn_, N, n, HC * N); + gm("mtp.fc_hidden.weight", XHn_, mh2_, N, n, (size_t) c * MAXT * N * 4); + } + { + mtl::Dispatch d; + d.p = P(lmtp_, "mtp_combine", {}, "mtp_combine"); + CombArgs a{N, HC, (uint32_t) n, (uint32_t) (MAXT * N)}; + d.args = {mtl::Arg::B(mh2_), mtl::Arg::B(me2_), mtl::Arg::B(mR_), mtl::Arg::V(a)}; + d.grid[0] = (n * HC * N + 255) / 256; d.tg[0] = 256; + push(ds, d); + } + // ---- the attention hyper-connection, K/V (and the indexer keys) into the layer's own cache + gr_read_m(ds, L + "attn_hyper_connection", mR_, mbo_, minj_[1], minj_[0], false, XA_, false, n); + gm(L + "self_attn.k_proj.weight", XA_, kcur_, 512, n); + nb_ = true; + gm(L + "self_attn.v_proj.weight", XA_, vcur_, 512, n); + nb_ = true; + gm(L + "self_attn.indexer.index_qk_proj.weight", XA_, midx_, 640, n); + if (full) { + nb_ = true; + gm(L + "self_attn.q_proj.weight", XA_, qfull_, 12288, n); + } + QsaArgs qa{(uint32_t) cell0, 0, (uint32_t) max_cells_, (uint32_t) n, eps_}; + { + mtl::Dispatch d; + d.p = P(lqsa_, "qsa_prep", qsa_c(), qsa_k()); + d.args = {mtl::Arg::B(qfull_), mtl::Arg::B(kcur_), mtl::Arg::B(vcur_), mtl::Arg::B(midx_), mtl::Arg::B(midx_, 512 * 4), + mtl::Arg::B(mnorm_, NM(L + "self_attn.q_norm.weight")), mtl::Arg::B(mnorm_, NM(L + "self_attn.k_norm.weight")), + mtl::Arg::B(mnorm_, NM(L + "self_attn.indexer.q_layernorm.weight")), + mtl::Arg::B(mnorm_, NM(L + "self_attn.indexer.k_layernorm.weight")), + mtl::Arg::B(cs_), mtl::Arg::B(qcur_), mtl::Arg::B(qidx_), mtl::Arg::B(mkq_[0]), mtl::Arg::B(mvq_[0]), + mtl::Arg::B(mksc_[0]), mtl::Arg::B(mvsc_[0]), mtl::Arg::B(mtail_[0]), mtl::Arg::B(mdead_[0]), + mtl::Arg::B(mpooled_[0]), mtl::Arg::V(qa), mtl::Arg::B(mrope_)}; + d.grid[0] = 33; d.grid[1] = n; d.tg[0] = 256; + push(ds, d); + } + if (!full) return; + // ---- dense attention over the window of cells up to this one (the draft layer's rule), the output gate, o_proj + { + int* ids = mqids_.as() + (size_t) slot * 2051; + const int64_t first = std::max(0, cell0 - 2050); + for (int64_t c = first; c <= cell0; ++c) ids[c - first] = (int) c; + mtl::Dispatch at; + at.p = P(lqsa_, "qsa_attn", {}, "qsa_attn"); + at.args = {mtl::Arg::B(qcur_), mtl::Arg::B(mkq_[0]), mtl::Arg::B(mvq_[0]), mtl::Arg::B(mksc_[0]), mtl::Arg::B(mvsc_[0]), + mtl::Arg::B(mqids_, (size_t) slot * 2051 * 4), mtl::Arg::B(part_), mtl::Arg::V(qa)}; + at.grid[0] = 33; at.grid[1] = 2; at.grid[2] = 1; at.tg[0] = 256; + push(ds, at); + mtl::Dispatch m; + m.p = P(lqsa_, "qsa_merge", {}, "qsa_merge"); + QsaArgs am = qa; + am.n_tok = MAXT; + m.args = {mtl::Arg::B(part_), mtl::Arg::B(qfull_), mtl::Arg::B(XQ_.b), mtl::Arg::V(am)}; + m.grid[0] = 24; m.grid[1] = 1; m.tg[0] = 256; + push(ds, m); + gm(L + "self_attn.o_proj.weight", XQ_, mbo_, N, 1); + } + // ---- the MLP hyper-connection (the attention write folded in), the MoE, the write + final mixer, the head + gr_read_m(ds, L + "mlp_hyper_connection", mR_, mbo_, minj_[0], minj_[1], true, XF_, false, 1); + gm(L + "mlp.gate.weight", XF_, logits_, 512, 1); + nb_ = true; + gm(L + "mlp.shared_expert.gate_proj.weight", XF_, gsh_, NFF, 1); + nb_ = true; + gm(L + "mlp.shared_expert.up_proj.weight", XF_, ush_, NFF, 1); + nb_ = true; + gm(L + "mlp.shared_expert_gate.weight", XF_, sg_, 1, 1); + { + MoeArgs a{}; + const gguf::Tensor &tgu = TM(L + "mlp.experts.gate_up_proj"), &td = TM(L + "mlp.experts.down_proj"); + a.gu_row = (uint32_t) ref::row_bytes(tgu.type, N); + a.d_row = (uint32_t) ref::row_bytes(td.type, NFF); + a.gate_off = tgu.offset; + a.up_off = tgu.offset + (uint64_t) NFF * a.gu_row; // [gate rows | up rows] per expert + a.down_off = td.offset; + a.gu_stride = (uint64_t) a.gu_row * 2 * NFF; + a.d_stride = (uint64_t) a.d_row * N; + a.n_embd = N; a.n_ff = NFF; a.n_tok = 1; a.topk = TOPK; a.n_expert = NEXP; + const bool dq2 = td.type == ref::Q2_0; + const int gl = 32, gr = 1, gs = 4, ds_ = 4; + // the down tile: Q2_0 8,1; else 4,2 - but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there: + // 24-36 GB/s, mac-kern 513fc615) + const bool m4 = !m5_layout(); + const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4, + dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2; + a.gx_gu = NFF / (gs * (32 / gl) * gr); + a.gx_dn = N / (ds_ * (32 / dl) * dr); + mtl::Dispatch r; + r.p = P(lib_, "moe_route", {}, "moe_route"); + r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)}; + r.tg[0] = 512; + push(ds, r); + { + GlueArgs ga{NFF, MAXT, N, TOPK}; + mtl::Dispatch d; + d.p = P(lib_, "swiglu_prep", {}, "swiglu_prep"); + d.args = {mtl::Arg::B(gsh_), mtl::Arg::B(ush_), mtl::Arg::B(XH_.b), mtl::Arg::V(ga)}; + d.grid[0] = NFF / 16; d.grid[1] = 1; d.tg[0] = 16; + nb_ = true; + push(ds, d); + } + const std::string gun = "moe_gu_" + std::to_string(tgu.type) + "_" + std::to_string(gl) + "_" + std::to_string(gr) + + "_" + std::to_string(gs); + const std::string dnn = "moe_down_" + std::to_string(dq2 ? 142 : td.type) + "_" + std::to_string(dl) + "_" + + std::to_string(dr) + "_" + std::to_string(ds_); + mtl::Dispatch g; + g.p = P(lib_, gun, mtl::Constants().b(23, metal_ctab()), gun + (metal_ctab() ? "#ct" : "")); + g.args = {mtl::Arg::B(Wm_), mtl::Arg::B(XF_.b, XF_.plane_bytes() * xkind(tgu.type)), mtl::Arg::B(h_), mtl::Arg::B(plan_), + mtl::Arg::V(a), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4)}; + g.tg[0] = 32 * gs; + g.indirect = &plan_; + g.ind_off = P_IGU * 4; + push(ds, g); + nb_ = true; + gm(L + "mlp.shared_expert.down_proj.weight", XH_, ysh_, N, 1); + mtl::Dispatch dd; + dd.p = P(lib_, dnn, {}, dnn); + dd.args = {mtl::Arg::B(Wm_), mtl::Arg::B(h_), mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(h_)}; + dd.tg[0] = 32 * ds_; + dd.indirect = &plan_; + dd.ind_off = P_IDN * 4; + push(ds, dd); + GlueArgs ca{N, 1, N, TOPK}; + mtl::Dispatch cb; + cb.p = P(lib_, "moe_combine_sh", {}, "moe_combine_sh"); + cb.args = {mtl::Arg::B(ye_), mtl::Arg::B(plan_), mtl::Arg::B(ysh_), mtl::Arg::B(sg_), mtl::Arg::B(mbo_), mtl::Arg::V(ca)}; + cb.grid[0] = (N + 255) / 256; cb.tg[0] = 256; + push(ds, cb); + } + // the layer's write (R updated in place: the next draft step's input residual) and the final mixer + gr_read_m(ds, "mtp.hyper_connection_mixer", mR_, mbo_, minj_[1], minj_[0], true, XA_, true, 1); + const int64_t nv = dvocab_.empty() ? n_vocab_ : (int64_t) dvocab_.size(); + if (dvocab_.empty()) gemv(ds, T("output.weight"), XA_, mlm_, (size_t) n_vocab_, 1); + else gemv(ds, mhead_t_, XA_, mlm_, (size_t) nv, 1, 0, &mhead_); + { + ArgArgs a{(uint32_t) nv, 120}; + mtl::Dispatch d1, d2; + d1.p = P(lsmp_, "argmax_part", {}, "argmax_part"); + d1.args = {mtl::Arg::B(mlm_), mtl::Arg::B(apv_), mtl::Arg::B(api_), mtl::Arg::V(a)}; + d1.grid[0] = 120; d1.grid[1] = 1; d1.tg[0] = 256; + push(ds, d1); + d2.p = P(lsmp_, "argmax_final", {}, "argmax_final"); + d2.args = {mtl::Arg::B(apv_), mtl::Arg::B(api_), mtl::Arg::B(tokout_), mtl::Arg::V(a)}; + d2.grid[0] = 1; d2.tg[0] = 256; + push(ds, d2); + } +} + +// the prompt's K/V: cell cell0 + i pairs the main window's final residual row i with the token at cell0 + i + 1 +bool Window::mtp_prefill(const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err) { + if (!mtp_on_) { err = "metal mtp: not loaded"; return false; } + mrope_sync(cell0, n); + // every group of MAXT rows in one command buffer: the embedding rows of all of them in mpemb_ (row c + i) + if ((int64_t) mpemb_rows_ < n) { + if (mpemb_rows_ > 0) cx_->free(mpemb_); + mpemb_rows_ = (size_t) std::max(n, 64); + mpemb_ = cx_->alloc(4 * mpemb_rows_ * N); + } + membp_ = &mpemb_; + std::vector ds; + nb_ = false; + for (int64_t c = 0; c < n; c += MAXT) { + const int k = (int) std::min(MAXT, n - c); + mtp_encode(ds, mpsrc_ ? *mpsrc_ : mtp_src_ ? *mtp_src_ : R_, (size_t) c * HC * N * 4, k, next_tokens + c, + cell0 + c, false, (int) c, 0); + } + membp_ = nullptr; + cx_->run_graph(ds); + return true; +} + +void Window::mtp_pick(std::vector& ds, int j, bool embed) { + struct PickArgs { uint32_t nv, j, use_map; }; + const int64_t nv = dvocab_.empty() ? n_vocab_ : (int64_t) dvocab_.size(); + mtl::Dispatch d; + d.p = P(lchain_, "mtp_pick", {}, "mtp_pick"); + PickArgs pa{(uint32_t) nv, (uint32_t) j, dvocab_.empty() ? 0u : 1u}; + d.args = {mtl::Arg::B(mlm_), mtl::Arg::B(tokout_), mtl::Arg::B(dmap_), mtl::Arg::B(dout_), mtl::Arg::B(dprob_), + mtl::Arg::V(pa)}; + d.tg[0] = 1024; + push(ds, d); + if (!embed) return; + const gguf::Tensor& te = T("token_embd.weight"); + struct EmbArgs { uint32_t n, row_bytes, j; } ea{(uint32_t) N, (uint32_t) ref::row_bytes(te.type, N), (uint32_t) j}; + mtl::Dispatch e; + const std::string fn = "mtp_embed_" + std::to_string(te.type); + e.p = P(lchain_, fn, {}, fn); + e.args = {mtl::Arg::B(W_, te.offset), mtl::Arg::B(dout_), mtl::Arg::B(posv_), mtl::Arg::B(memb_), mtl::Arg::V(ea)}; + e.tg[0] = 256; + push(ds, e); +} + +// One round (mtp.hpp's draft()): the window's rows 0..a-1 caught up (K/V), then the full layer on row a (draft 0) and +// on its own output for the next drafts. By default the whole round is one command buffer (the pick, its +// probability and the next step's embedding on the GPU); STRATA_METAL_MTP_CHAIN=0: a command buffer per step, the +// host picking and dequantizing. The embedding rows: the catch-up's at memb_ rows 0..a-1, draft 0's at row a. +bool Window::mtp_draft(int Tn, const int32_t* tokens, int64_t p, int a, int max_steps, int32_t* drafts, float* probs, + float min_p, int* n_drafts, std::string& err) { + if (!mtp_on_) { err = "metal mtp: not loaded"; return false; } + (void) Tn; + mrope_sync(p, a + max_steps + 1); + int n = 0; + int32_t tok = tokens[a]; + if (chain_ && max_steps > 0) { + std::vector ds; + nb_ = false; + if (a > 0) mtp_encode(ds, R_, 0, a, tokens, p, false); + mtp_encode(ds, R_, (size_t) a * HC * N * 4, 1, &tok, p + a, true, a, 0); + mtp_pick(ds, 0, max_steps > 1); + for (int j = 1; j < max_steps; ++j) { + mtp_encode(ds, mR_, 0, 1, nullptr, p + a + j, true, 0, j); + mtp_pick(ds, j, j + 1 < max_steps); + } + cx_->run_graph(ds); + for (int j = 0; j < max_steps; ++j) { + const float pj = dprob_.as()[j]; + drafts[j] = dout_.as()[j]; + if (probs) probs[j] = pj; + ++n; + if (min_p > 0.0f && pj < min_p) break; + } + if (n_drafts) *n_drafts = n; + return true; + } + for (int j = 0; j < max_steps; ++j) { + std::vector ds; + nb_ = false; + if (j == 0) { + if (a > 0) mtp_encode(ds, R_, 0, a, tokens, p, false); + mtp_encode(ds, R_, (size_t) a * HC * N * 4, 1, &tok, p + a, true, a, 0); + } else { + mtp_encode(ds, mR_, 0, 1, &tok, p + a + j, true); + } + cx_->run_graph(ds); + const int32_t di = tokout_.as()[0]; + const int32_t d = dvocab_.empty() ? di : dvocab_[(size_t) di]; + // the draft's probability: softmax of the draft head's logits at the pick + const float* lg = mlm_.as(); + const float mx = lg[di]; + const int64_t nv = dvocab_.empty() ? n_vocab_ : (int64_t) dvocab_.size(); + double s = 0.0; + for (int64_t i = 0; i < nv; ++i) s += std::exp((double) (lg[i] - mx)); + const float pj = (float) (1.0 / s); + drafts[j] = d; + if (probs) probs[j] = pj; + tok = d; + ++n; + if (min_p > 0.0f && pj < min_p) break; + } + if (n_drafts) *n_drafts = n; + return true; +} + +bool Window::copy_logits(int t, float* host) const { + if (t < 0 || t >= last_T_) return false; + std::memcpy(host, lm_.as() + (size_t) t * n_vocab_, (size_t) n_vocab_ * 4); + return true; +} + + +// ================================================================================== prompts through the matrix units +// A prompt chunk of T rows (mac-kern's tests/prompt_chain.cpp, on the engine's state): every dense matrix dequantized +// to fp16 once a chunk and multiplied on the matrix units (MPP), the GDN scan, QSA over the chunk, the routed experts +// as grouped GEMMs (v2 on the M4, v3 on the M5). The GDN state, the conv history, the K/V and indexer cells and the +// PLE history are the decode window's own, so decoding goes on from them. With SSD streaming the experts of a layer +// are put in place at the layer's handshake (the arena's misses) and the next layer's are streamed in behind it. +// STRATA_METAL_PROMPT_GEMM=0: prompts through the decode window (8 rows at a time). Its numerics differ from the +// window's (fp16 operands): STRATA_METAL_PROMPT_GEMM_MIN (default 64) keeps shorter prompts on the window path. +namespace { +struct PgDq { uint32_t n_in, n_out, row_bytes, perm; }; +struct PgGemm { int M, N, K, pad; }; +struct PgEpi { uint32_t n_tok, n_embd, hc, hc_lr; }; +struct PgGdnP { uint32_t n_tok, h_k, h_v, channels; float eps; uint32_t f32; }; +struct PgPlan { uint32_t n_ent, topk, n_expert, tm, n_blocks; }; +struct PgMm { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff; }; +int pg_unit(int t) { return t == ref::Q2_0 || t == ref::IQ3_S || t == ref::BF16 || t == ref::F32 ? 16 : 32; } +// the v2 (M4) expert tiles: the K-quants need TK 128, Q8_0 / IQ4_XS TK 32 (their 32-value units), the rest TK 16 +std::string gu2_tile(int t) { + if (t == ref::Q4_K || t == ref::Q5_K || t == ref::Q6_K) return "_16_16_128_2"; // the M4's best (513fc615) + if (t == ref::Q8_0 || t == ref::IQ4_XS) return "_16_64_32_2"; + return "_16_64_16_2"; +} +std::string dn2_tile(int t) { return t == ref::IQ4_NL || t == ref::Q5_1 || t == ref::Q8_0 ? "_16_64_32_2" : "_16_64_16_2"; } +} // namespace + +struct Window::PromptGemm { + int MT = 0, QS = 256, TM = 16, maxrows = 0, nblk_max = 0; + int lib = -1, lgdn = -1, lqsa = -1; + bool k4 = false, gu_v3 = false, dn_v3 = false, flash = false; + mtl::Buffer R, bo, inj[2], xnh, V, loh, mixed, mixh, qkv, z, al, be, ygh, gqn, gkn, gvc, ggb, gkq, go, yg, qfull, kcur, vcur, + iqr, ikr, qcur, qcurh, qidx, scores, qids, part, XQ, xqh, logits, ids, wts, bcnt, boff, eoff, tiles, meta, stok, sw, pos, + XS, HS, YE, yr, gsh, ush, hh, ys, sgv, emb, embh, pkey, pval, keyn, qn, gated, normW; + mtl::Buffer WD, WU, Wqkv, Wz, Wal, Wbe, Wout, Wq, Wk, Wv, Wiq, Wik, Wr, Wsg, Wshg, Wshu, Wshd, Wpk, Wpv; + mtl::Buffer mixr; // Gyro: the rotated gate/up input rows [MT][N] +}; + +int Window::prompt_chunk_rows() const { return pg_ ? pg_->MT : 0; } + +bool Window::prompt_gemm_on(int64_t n) const { + static const bool on = [] { + const char* e = std::getenv("STRATA_METAL_PROMPT_GEMM"); + return e ? std::string(e) != "0" : true; + }(); + static const int64_t mn = [] { + const char* e = std::getenv("STRATA_METAL_PROMPT_GEMM_MIN"); + return e ? (int64_t) std::atoll(e) : (int64_t) 64; + }(); + return on && n >= mn && pg_ != nullptr; +} + +bool Window::prompt_gemm_init(std::string& err) { + if (pg_) return true; + if (stream_ && !st_->ready()) return stream_arena(err); // it makes the prompt buffers after the arena + const char* ec = std::getenv("STRATA_METAL_PROMPT_CHUNK"); + // Streaming: every chunk makes one pass over the layer's experts that are not in the arena, so a longer chunk + // streams each expert fewer times. M5 Max, UD-Q5, warm: 1024 -> 700 tok/s at 4K, 2048 -> 920, 4096 -> 1,100 (32K the + // same). The longer chunk's buffers (4.3 GB) come out of the arena only while a prompt runs, so it is the default on a + // big Mac; a small one keeps 1024 (the buffers would cost it cache). + uint64_t phys = 0; + size_t plen = sizeof phys; + sysctlbyname("hw.memsize", &phys, &plen, nullptr, 0); + const int stream_chunk = phys >= (96ull << 30) ? 4096 : 1024; + const int C = std::max(64, ec ? std::atoi(ec) : (stream_ ? stream_chunk : 2048)) / 64 * 64; + auto pg = std::make_unique(); + PromptGemm& p = *pg; + p.MT = C; + p.QS = std::min(C, 256); + p.lib = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", "gemv.metal", + "prompt.metal"}); + p.lgdn = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + p.lqsa = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + p.k4 = ::coop_k4(*cx_); + const char* eg = std::getenv("STRATA_METAL_PROMPT_EXPERTS"); // v2 | v3 (default: v3 with the M5's layout) + p.gu_v3 = p.dn_v3 = eg ? std::string(eg) == "v3" : p.k4; + if (T("blk.0.ffn_gate_exps.weight").type >= 144) p.gu_v3 = p.dn_v3 = false; // the trellis experts have the v2 kernels only + const char* eq = std::getenv("STRATA_METAL_PROMPT_QSA"); // v1 | flash (default: flash with the M5's layout) + p.flash = eq ? std::string(eq) == "flash" : p.k4; + const size_t MT = (size_t) C; + p.maxrows = C * TOPK + NEXP * p.TM; + p.nblk_max = (C * TOPK + 1023) / 1024; + const int MB = (int) max_cells_ / 4 + 2; + // the buffers: measured first, then placed - on a streaming Mac inside the expert arena (a region of its slots + // reserved only while a prompt runs: the prompt's scratch costs no decode memory), else allocated + uint64_t need = 0; + uint8_t* region = nullptr; + bool measure = true; + auto AB = [&](size_t bytes) -> mtl::Buffer { + const size_t b = (bytes + 16383) / 16384 * 16384; + if (measure) { need += b; return mtl::Buffer{}; } + if (region == nullptr) return cx_->alloc(bytes); + mtl::Buffer r = cx_->wrap(region, b); + region += b; + return r; + }; + auto F32 = [&](size_t n) { return AB(4 * n); }; + auto F16 = [&](size_t n) { return AB(2 * n); }; + auto make = [&] { + const size_t mr = (size_t) p.maxrows; + p.R = F32(MT * HC * N); p.bo = F32(MT * N); p.inj[0] = F32(MT * HC); p.inj[1] = F32(MT * HC); p.xnh = F16(MT * HC * N); + p.V = F32(MT * (LR + HC)); p.loh = F16(MT * LR); p.mixed = F32(MT * N); p.mixh = F16(MT * N); + p.qkv = F32(MT * 10240); p.z = F32(MT * 6144); p.al = F32(MT * 48); p.be = F32(MT * 48); p.yg = F32(MT * 6144); + p.ygh = F16(MT * 6144); p.gqn = F32(MT * 16 * 128); p.gkn = F32(MT * 16 * 128); p.gvc = F32(MT * 6144); + p.ggb = F32(MT * 96); p.gkq = F32(MT * 16); p.go = F32(MT * 6144); + p.qfull = F32(MT * 12288); p.kcur = F32(MT * 512); p.vcur = F32(MT * 512); p.iqr = F32(MT * 512); p.ikr = F32(MT * 128); + p.qcur = F32(MT * 24 * 256); p.qcurh = F16(MT * 24 * 256); p.qidx = F32(MT * 4 * 128); p.scores = F32(MT * (size_t) MB); + p.qids = AB(4 * MT * 2051); + p.part = p.flash ? F32(16) : F32((size_t) p.QS * 33 * 24 * 258); + p.XQ = F32(4 * MT * 6144 + 3 * (size_t) p.QS * 6144 + MT * 384); p.xqh = F16(MT * 6144); + p.logits = F32(MT * NEXP); p.ids = AB(4 * MT * TOPK); p.wts = F32(MT * TOPK); + p.bcnt = AB(4 * (size_t) p.nblk_max * NEXP); p.boff = AB(4 * (size_t) p.nblk_max * NEXP); + p.eoff = AB(4 * NEXP); p.tiles = AB(8 * (mr / p.TM + 1)); p.meta = AB(64); + p.stok = AB(4 * mr); p.sw = F32(mr); p.pos = AB(4 * MT * TOPK); p.XS = F16(mr * N); p.HS = F16(mr * NFF); + p.YE = F32(mr * N); p.yr = F32(MT * N); p.gsh = F32(MT * NFF); p.ush = F32(MT * NFF); p.hh = F16(MT * NFF); + p.ys = F32(MT * N); p.sgv = F32(MT); + if (had_on_) p.mixr = F32(MT * N); + p.WD = F16((size_t) (LR + HC) * HC * N); p.WU = F16((size_t) HC * N * LR); p.Wqkv = F16((size_t) 10240 * N); + p.Wz = F16((size_t) 6144 * N); p.Wal = F16((size_t) 48 * N); p.Wbe = F16((size_t) 48 * N); p.Wout = F16((size_t) N * 6144); + p.Wq = F16((size_t) 12288 * N); p.Wk = F16((size_t) 512 * N); p.Wv = F16((size_t) 512 * N); p.Wiq = F16((size_t) 512 * N); + p.Wik = F16((size_t) 128 * N); p.Wr = F16((size_t) NEXP * N); p.Wsg = F16((size_t) N); p.Wshg = F16((size_t) NFF * N); + p.Wshu = F16((size_t) NFF * N); p.Wshd = F16((size_t) N * NFF); + if (ple_on_) { + p.emb = F32(MT * N); p.embh = F16(MT * N); p.pkey = F32(MT * 10240); p.pval = F32(MT * N); p.keyn = F32(MT * 10240); + p.qn = F32(MT * 10240); p.gated = F32(MT * 10240); p.normW = F32(MT * 10240); + p.Wpk = F16((size_t) 10240 * N); p.Wpv = F16((size_t) N * N); + } + }; + make(); + measure = false; + if (stream_) { + region = st_->region(need); + if (region == nullptr) { err = "metal prompt: no arena region for the prompt buffers"; return false; } + } + make(); + pg_ = pg.release(); + std::fprintf(stderr, "strata metal: prompts on the matrix units in chunks of %d (experts %s, QSA %s, matmul2d layout %s; " + "%.2f GB of buffers%s)\n", C, p.gu_v3 ? "v3" : "v2", p.flash ? "flash" : "v1", p.k4 ? "K4" : "other", + need / 1e9, stream_ ? ", in the expert arena while a prompt runs" : ""); + (void) err; + return true; +} + +bool Window::prompt_chunk(const int32_t* tokens, int Tn, int64_t pos0, std::string& err) { + PromptGemm& p = *pg_; + p.meta.as()[2] = p.gu_v3 ? NFF / (32 * 4) : NFF / 64; + p.meta.as()[5] = p.dn_v3 ? N / (64 * 2) : N / 64; + auto TN = [this](const std::string& n) -> const gguf::Tensor& { return this->T(n); }; + if (Tn < 1 || Tn > p.MT) { err = "metal prompt: chunk size out of range"; return false; } + if (pos0 + Tn + 1 >= max_cells_) { err = "metal prompt: the chunk runs past the context"; return false; } + if (stream_ && !st_->ready() && !stream_arena(err)) return false; + if (pend_fold_ > 0) { // the last decode window's commit, still pending + std::vector fd; + nb_ = false; + fold(fd, pend_fold_, 1 - hp_); + cx_->run_graph(fd); + pend_fold_ = 0; + } + const int T = Tn; + // ---- the embeddings into every stream, the PLE rows + { + const gguf::Tensor& te = TN("token_embd.weight"); + const size_t rb = ref::row_bytes(te.type, N); + mrope_sync(pos0, T); + for (int t = 0; t < T; ++t) { + float* r = p.R.as() + (size_t) t * HC * N; + if (const float* img = img_row(pos0 + t)) std::memcpy(r, img, (size_t) N * 4); + else ref::dequant(te.type, (const uint8_t*) W_.ptr + te.offset + (size_t) tokens[t] * rb, r, N); + for (int c = 1; c < HC; ++c) std::memcpy(r + (size_t) c * N, r, (size_t) N * 4); + } + } + strata::core::SessionState& ss = *ss_; + if (ple_on_) { + std::vector rows((size_t) T * strata::kernels::PLE_N_HEADS); + int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]}; + for (int t = 0; t < T; ++t) { + strata::kernels::ngram_rows(&tokens[t], prev, 1, ss.ple.consts, rows.data() + (size_t) t * strata::kernels::PLE_N_HEADS); + prev[0] = prev[1]; + prev[1] = tokens[t]; + } + if (!ss.ple.table->gather_batch(rows.data(), (size_t) T, p.emb.as(), err)) return false; + } + std::vector ds; + std::vector cuts; + nb_ = false; + auto D = [&](const mtl::Pipeline* pp, std::vector a, uint32_t gx, uint32_t gy, uint32_t tg, uint32_t gz = 1) { + mtl::Dispatch d; + d.p = pp; + d.args = std::move(a); + d.grid[0] = gx; d.grid[1] = gy; d.grid[2] = gz; + d.tg[0] = tg; + push(ds, d); + }; + auto PP = [&](int lib, const std::string& fn, const mtl::Constants& c = {}, const std::string& key = "") { + return P(lib, fn, c, "pg:" + (key.empty() ? fn : key)); + }; + auto gemm = [&](const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Y, int M, int Nn, int K) { + const int tm = 64, tn = Nn <= 64 ? 32 : 64; + const std::string fn = "gemm_f16_" + std::to_string(tm) + "_" + std::to_string(tn) + "_4"; + PgGemm g{M, Nn, K, 0}; + D(PP(p.lib, fn), {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, (Nn + tn - 1) / tn, (M + tm - 1) / tm, 128); + }; + auto dqb = [&](const mtl::Buffer& src, size_t off, int type, int n_in, int n_out, const mtl::Buffer& out, size_t out_off = 0, + uint32_t perm = 0) { + PgDq a{(uint32_t) n_in, (uint32_t) n_out, (uint32_t) ref::row_bytes(type, n_in), perm}; + D(PP(p.lib, "dq_f16_" + std::to_string(type)), {mtl::Arg::B(src, off), mtl::Arg::B(out, out_off), mtl::Arg::V(a)}, + (n_in / pg_unit(type) + 63) / 64, n_out, 64); + }; + auto dq = [&](const gguf::Tensor& t, const mtl::Buffer& out, size_t out_off = 0, uint32_t perm = 0) { + const int n_in = (int) t.shape[0], n_out = t.shape.size() > 1 ? (int) t.shape[1] : 1; + if (hc_q8_ && q8off_.count(t.name)) dqb(Q8_, q8off_.at(t.name), ref::Q8_0, n_in, n_out, out, out_off, perm); + else dqb(W_, t.offset, t.type, n_in, n_out, out, out_off, perm); + }; + auto act16 = [&](const mtl::Buffer& x, const mtl::Buffer& xh, size_t n) { + const uint32_t nn = (uint32_t) n; + D(PP(p.lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 255) / 256, 1, 256); + }; + int ci = 0; + auto hc_read = [&](const std::string& pre, bool apply) { + const gguf::Tensor &tn = TN(pre + "_norm.weight"), &tdn = TN(pre + "_down.weight"), &tin = TN(pre + "_inject.weight"), + &tup = TN(pre + "_up.weight"); + dq(tdn, p.WD); + dq(tin, p.WD, (size_t) 2 * LR * HC * N); + dq(tup, p.WU, 0, HC); + GrArgs a{N, HC, LR, (uint32_t) T, apply ? 1u : 0u, eps_}; + PgEpi ea{(uint32_t) T, N, HC, LR}; + D(PP(p.lib, "gr_norm_h"), {mtl::Arg::B(p.R), mtl::Arg::B(p.bo), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), + mtl::Arg::B(p.xnh), mtl::Arg::V(a)}, HC, T, 256); + gemm(p.xnh, p.WD, p.V, T, LR + HC, HC * N); + D(PP(p.lib, "gr_epi_down"), {mtl::Arg::B(p.V), mtl::Arg::B(p.loh), mtl::Arg::B(p.inj[1 - ci]), mtl::Arg::V(ea)}, + (T * (LR + HC) + 255) / 256, 1, 256); + PgGemm g{T, HC * N, LR, 0}; + D(PP(p.lib, "gemm_hc_up_64_64_4"), {mtl::Arg::B(p.loh), mtl::Arg::B(p.WU), mtl::Arg::B(p.xnh), mtl::Arg::B(p.mixed), + mtl::Arg::B(p.mixh), mtl::Arg::V(g), mtl::Arg::V(ea)}, HC * N / 64, (T + 63) / 64, 128); + ci = 1 - ci; + }; + // the PLE projections (the window's gemv pair as GEMMs) + if (ple_on_) { + act16(p.emb, p.embh, (size_t) T * N); + dq(TN("blk.1.ple_key.weight"), p.Wpk); + dq(TN("blk.1.ple_value.weight"), p.Wpv); + gemm(p.embh, p.Wpk, p.pkey, T, 10240, N); + gemm(p.embh, p.Wpv, p.pval, T, N, N); + } + const int MB = (int) max_cells_ / 4 + 2; + (void) MB; + for (int l = 0; l < 48; ++l) { + const std::string b = "blk." + std::to_string(l) + "."; + hc_read(b + "hc_attn", l > 0 && !(l == 1 && ple_on_)); + if (l % 4 != 3) { + const int gi = l - l / 4; + dq(TN(b + "attn_qkv.weight"), p.Wqkv); + dq(TN(b + "attn_gate.weight"), p.Wz); + dq(TN(b + "ssm_alpha.weight"), p.Wal); + dq(TN(b + "ssm_beta.weight"), p.Wbe); + dq(TN(b + "ssm_out.weight"), p.Wout); + gemm(p.mixh, p.Wqkv, p.qkv, T, 10240, N); + gemm(p.mixh, p.Wz, p.z, T, 6144, N); + gemm(p.mixh, p.Wal, p.al, T, 48, N); + gemm(p.mixh, p.Wbe, p.be, T, 48, N); + PgGdnP ga{(uint32_t) T, 16, 48, 10240, eps_, 0}; + D(PP(p.lgdn, "gdn_pre"), + {mtl::Arg::B(p.qkv), mtl::Arg::B(ghist_[hp_][(size_t) gi]), mtl::Arg::B(W_, TN(b + "ssm_conv1d.weight").offset), + mtl::Arg::B(p.al), mtl::Arg::B(p.be), mtl::Arg::B(W_, TN(b + "ssm_dt.bias").offset), mtl::Arg::B(W_, TN(b + "ssm_a").offset), + mtl::Arg::B(p.gqn), mtl::Arg::B(p.gkn), mtl::Arg::B(p.gvc), mtl::Arg::B(p.ggb), mtl::Arg::B(p.gkq), + mtl::Arg::B(ghist_[1 - hp_][(size_t) gi]), mtl::Arg::V(ga)}, 16 + 24, T, 256); + D(PP(p.lgdn, gdn_scan_name(cx_->device_name(), T), mtl::Constants().i(6, 0), gdn_scan_name(cx_->device_name(), T)), + {mtl::Arg::B(gst_[(size_t) gi]), mtl::Arg::B(p.gqn), mtl::Arg::B(p.gkn), mtl::Arg::B(p.gvc), mtl::Arg::B(p.ggb), + mtl::Arg::B(p.gkq), mtl::Arg::B(p.go), mtl::Arg::V(ga)}, 48 * 4, 1, 32 * 4); + D(PP(p.lgdn, "gdn_post"), {mtl::Arg::B(p.go), mtl::Arg::B(p.z), mtl::Arg::B(W_, TN(b + "ssm_norm.weight").offset), + mtl::Arg::B(p.ygh), mtl::Arg::B(p.yg), mtl::Arg::V(ga)}, 48, T, 128); + gemm(p.ygh, p.Wout, p.bo, T, N, 6144); + } else { + const int qi = l / 4; + dq(TN(b + "attn_q.weight"), p.Wq); + dq(TN(b + "attn_k.weight"), p.Wk); + dq(TN(b + "attn_v.weight"), p.Wv); + dq(TN(b + "indexer.q_proj.weight"), p.Wiq); + dq(TN(b + "indexer.k_proj.weight"), p.Wik); + dq(TN(b + "attn_output.weight"), p.Wout); + gemm(p.mixh, p.Wq, p.qfull, T, 12288, N); + gemm(p.mixh, p.Wk, p.kcur, T, 512, N); + gemm(p.mixh, p.Wv, p.vcur, T, 512, N); + gemm(p.mixh, p.Wiq, p.iqr, T, 512, N); + gemm(p.mixh, p.Wik, p.ikr, T, 128, N); + QsaArgs a{(uint32_t) pos0, 0, (uint32_t) max_cells_, (uint32_t) T, eps_}; + D(PP(p.lqsa, "qsa_prep", qsa_c(), qsa_k()), + {mtl::Arg::B(p.qfull), mtl::Arg::B(p.kcur), mtl::Arg::B(p.vcur), mtl::Arg::B(p.iqr), mtl::Arg::B(p.ikr), + mtl::Arg::B(W_, TN(b + "attn_q_norm.weight").offset), mtl::Arg::B(W_, TN(b + "attn_k_norm.weight").offset), + mtl::Arg::B(W_, TN(b + "indexer.q_norm.weight").offset), mtl::Arg::B(W_, TN(b + "indexer.k_norm.weight").offset), + mtl::Arg::B(cs_), mtl::Arg::B(p.qcur), mtl::Arg::B(p.qidx), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), + mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(tail_[(size_t) qi]), + mtl::Arg::B(dead_[(size_t) qi]), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::V(a), mtl::Arg::B(mrope_)}, 33, T, 256); + if (pos0 + T > 2051) + D(PP(p.lqsa, "qsa_score_mm"), {mtl::Arg::B(p.qidx), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), + mtl::Arg::B(p.scores), mtl::Arg::V(a)}, + (uint32_t) (((pos0 + T) / 4 + 1 + 63) / 64), (T + 15) / 16, 128); + D(PP(p.lqsa, "qsa_topk"), {mtl::Arg::B(p.scores), mtl::Arg::B(p.qids), mtl::Arg::V(a)}, T, 1, 1024); + if (p.flash) { + const uint32_t nq = (uint32_t) T * 24 * 256; + D(PP(p.lqsa, "qsa_q_f16"), {mtl::Arg::B(p.qcur), mtl::Arg::B(p.qcurh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 256); + D(PP(p.lqsa, "qsa_flash", mtl::Constants().b(7, p.k4), "qsa_flash"), + {mtl::Arg::B(p.qcurh), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), mtl::Arg::B(ksc_[(size_t) qi]), + mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(p.qids), mtl::Arg::B(p.qfull), mtl::Arg::B(p.xqh), mtl::Arg::V(a)}, + 2, T, 128); + } else { + // v1 attention + merge in slices of QS rows (the partials of a slice only) + for (int t0 = 0; t0 < T; t0 += p.QS) { + const int ns = std::min(p.QS, T - t0); + QsaArgs as{(uint32_t) (pos0 + t0), 0, (uint32_t) max_cells_, (uint32_t) ns, eps_}; + D(PP(p.lqsa, "qsa_attn"), {mtl::Arg::B(p.qcur, (size_t) t0 * 24 * 256 * 4), mtl::Arg::B(kq_[(size_t) qi]), + mtl::Arg::B(vq_[(size_t) qi]), mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]), + mtl::Arg::B(p.qids, (size_t) t0 * 2051 * 4), mtl::Arg::B(p.part), mtl::Arg::V(as)}, + 33, 2, 256, (uint32_t) ns); + D(PP(p.lqsa, "qsa_merge"), {mtl::Arg::B(p.part), mtl::Arg::B(p.qfull, (size_t) t0 * 12288 * 4), + mtl::Arg::B(p.XQ, (size_t) t0 * 6144 * 4), mtl::Arg::V(as)}, 24, ns, 256); + } + act16(p.XQ, p.xqh, (size_t) T * 6144); + } + gemm(p.xqh, p.Wout, p.bo, T, N, 6144); + } + if (l == 0 && ple_on_) { + // layer 0's attention write applied to R happens in the FFN read below; the PLE block comes after the FFN + } + // ---- the MoE + hc_read(b + "hc_ffn", true); + const gguf::Tensor &tr = TN(b + "ffn_gate_inp.weight"), &tsg = TN(b + "ffn_gate_inp_shexp.weight"), + &tg_ = TN(b + "ffn_gate_shexp.weight"), &tu_ = TN(b + "ffn_up_shexp.weight"), + &td_ = TN(b + "ffn_down_shexp.weight"), &eg = TN(b + "ffn_gate_exps.weight"), + &ed = TN(b + "ffn_down_exps.weight"); + dq(tr, p.Wr); + dq(tsg, p.Wsg); + dq(tg_, p.Wshg); + dq(tu_, p.Wshu); + dq(td_, p.Wshd); + gemm(p.mixh, p.Wr, p.logits, T, NEXP, N); + const uint32_t nt = (uint32_t) T; + D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1, 256); + const int nent = T * TOPK, nblk = (nent + 1023) / 1024; + PgPlan pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) p.TM, (uint32_t) nblk}; + const uint32_t nclr = (uint32_t) p.maxrows; + D(PP(p.lib, "mm_clear"), {mtl::Arg::B(p.stok), mtl::Arg::V(nclr)}, (p.maxrows + 255) / 256, 1, 256); + D(PP(p.lib, "mm_hist"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024); + D(PP(p.lib, "mm_scan"), {mtl::Arg::B(p.bcnt), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.tiles), + mtl::Arg::B(p.meta), mtl::Arg::V(pa)}, 1, 1, NEXP); + D(PP(p.lib, "mm_place"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), + mtl::Arg::B(p.stok), mtl::Arg::B(p.sw), mtl::Arg::B(p.pos), mtl::Arg::V(pa)}, nblk, 1, 1024); + PgMm ma{eg.offset, TN(b + "ffn_up_exps.weight").offset, ed.offset, 0, 0, (uint32_t) ref::row_bytes(eg.type, N), + (uint32_t) ref::row_bytes(ed.type, NFF), N, NFF}; + ma.gu_stride = (uint64_t) ma.gu_row * NFF; + ma.d_stride = (uint64_t) ma.d_row * N; + if (stream_) { ma.gate_off = 0; ma.up_off = st_->up_in_slot(l); ma.down_off = st_->down_in_slot(l); } + const size_t ebase = moe_rebase(ma.gate_off, ma.up_off, ma.down_off); + const uint32_t ne_ = N; + const bool hg = had_gu(l), hd = had_dn(l); + if (hg) had_rows(ds, p.lib, "pg:", p.mixed, 0, p.mixr, 0, T, N, hsx_.handle ? &hsx_ : nullptr, false); + { + mtl::Dispatch g; + g.p = PP(p.lib, "mm_gather"); + g.args = {mtl::Arg::B(hg ? p.mixr : p.mixed), mtl::Arg::B(p.stok), mtl::Arg::B(p.XS), mtl::Arg::V(ne_)}; + g.tg[0] = 256; + g.indirect = &p.meta; + g.ind_off = 8 * 4; + push(ds, g); + } + const mtl::Buffer& EW = stream_ ? st_->chunk(0) : W_; + auto arena_args = [&](mtl::Dispatch& d) { + if (!stream_) return; + d.args.push_back(mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4)); + for (int c = 1; c < 4; ++c) d.args.push_back(mtl::Arg::B(st_->chunk(c))); + }; + const std::string ak = stream_ ? "#arena" : ""; + mtl::Dispatch gu; + if (p.gu_v3) { + const std::string fn = "mm_gu3_" + std::to_string(eg.type) + "_16_32_32_4"; + mtl::Constants c = mtl::Constants().b(7, p.k4); + if (stream_) c.b(20, true); + gu.p = PP(p.lib, fn, c, fn + ak); + gu.tg[0] = 32 * 4; + } else { + const std::string fn = "mm_gu2_" + std::to_string(eg.type) + gu2_tile(eg.type); + mtl::Constants c = mtl::Constants().i(5, 0); + if (stream_) c.b(20, true); + gu.p = PP(p.lib, fn, c, fn + ak); + gu.tg[0] = 64; + } + gu.args = {mtl::Arg::B(EW, ebase), mtl::Arg::B(p.XS), mtl::Arg::B(p.HS), mtl::Arg::B(p.tiles), mtl::Arg::V(ma)}; + arena_args(gu); + gu.indirect = &p.meta; + gu.ind_off = 2 * 4; + if (stream_) cuts.push_back((int) ds.size()); // the host's turn: this layer's experts put in place + push(ds, gu); + if (hd) had_rows_ind(ds, p.lib, "pg:", p.HS, p.meta); + mtl::Dispatch dn; + if (p.dn_v3) { + const std::string fn = "mm_down3_" + std::to_string(ed.type) + "_16_64_32_2"; + mtl::Constants c = mtl::Constants().b(7, p.k4); + if (stream_) c.b(20, true); + dn.p = PP(p.lib, fn, c, fn + ak); + dn.tg[0] = 32 * 2; + } else { + const std::string fn = "mm_down2_" + std::to_string(ed.type) + dn2_tile(ed.type); + mtl::Constants c; + if (stream_) c.b(20, true); + dn.p = PP(p.lib, fn, c, fn + ak); + dn.tg[0] = 64; + } + dn.args = {mtl::Arg::B(EW, ebase), mtl::Arg::B(p.HS), mtl::Arg::B(p.YE), mtl::Arg::B(p.tiles), mtl::Arg::V(ma)}; + arena_args(dn); + dn.indirect = &p.meta; + dn.ind_off = 5 * 4; + push(ds, dn); + const uint32_t ca[4] = {(uint32_t) T, N, TOPK, 0}; + D(PP(p.lib, "mm_combine"), {mtl::Arg::B(p.YE), mtl::Arg::B(p.pos), mtl::Arg::B(p.sw), mtl::Arg::B(p.yr), mtl::Arg::V(ca)}, + (T * N + 255) / 256, 1, 256); + gemm(p.mixh, p.Wshg, p.gsh, T, NFF, N); + gemm(p.mixh, p.Wshu, p.ush, T, NFF, N); + const uint32_t nh = (uint32_t) T * NFF; + D(PP(p.lib, "swiglu_h"), {mtl::Arg::B(p.gsh), mtl::Arg::B(p.ush), mtl::Arg::B(p.hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256); + gemm(p.hh, p.Wshd, p.ys, T, N, NFF); + gemm(p.mixh, p.Wsg, p.sgv, T, 1, N); + const uint32_t oa[2] = {(uint32_t) T, N}; + D(PP(p.lib, "moe_out"), {mtl::Arg::B(p.yr), mtl::Arg::B(p.ys), mtl::Arg::B(p.sgv), mtl::Arg::B(p.bo), mtl::Arg::V(oa)}, + (T * N + 255) / 256, 1, 256); + if (l == 0 && ple_on_) { + // layer 0's FFN write applied to R, then the PLE block (as the window, for T rows) + const gguf::Tensor& tn = TN("blk.1.hc_attn_norm.weight"); + GrArgs a{N, HC, LR, (uint32_t) T, 1u, eps_}; + D(P(lib_, "gr_norm", {}, "gr_norm"), {mtl::Arg::B(p.R), mtl::Arg::B(p.bo), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), + mtl::Arg::B(p.YE), mtl::Arg::V(a)}, HC, T, 256); + struct PleArgs { uint32_t n_embd, hc; float eps; } pl{N, HC, strata::kernels::NG_RMS_EPS}; + auto off = [&](const char* n) { return TN(n).offset; }; + D(P(lple_, "ple_gnorm", {}, "ple_gnorm"), {mtl::Arg::B(p.pkey), mtl::Arg::B(W_, off("blk.1.ple_norm_key.weight")), + mtl::Arg::B(p.keyn), mtl::Arg::V(pl)}, HC, T, 256); + nb_ = true; + D(P(lple_, "ple_gnorm", {}, "ple_gnorm"), {mtl::Arg::B(p.R), mtl::Arg::B(W_, off("blk.1.ple_norm_query.weight")), + mtl::Arg::B(p.qn), mtl::Arg::V(pl)}, HC, T, 256); + D(P(lple_, "ple_gate_norm", {}, "ple_gate_norm"), + {mtl::Arg::B(p.keyn), mtl::Arg::B(p.qn), mtl::Arg::B(p.pval), mtl::Arg::B(W_, off("blk.1.ple_norm_conv.weight")), + mtl::Arg::B(p.gated), mtl::Arg::B(p.normW), mtl::Arg::V(pl)}, HC, T, 256); + D(P(lple_, "ple_conv_add", {}, "ple_conv_add"), + {mtl::Arg::B(pleHist_), mtl::Arg::B(p.normW), kw_arg(), mtl::Arg::B(p.gated), + mtl::Arg::B(p.R), mtl::Arg::V(pl)}, HC * N / 256, T, 256); + } + } + // the last layer's write applied to R (the draft layer's K/V read the final residual rows) + { + const gguf::Tensor& tn = TN("output_hc_norm.weight"); + GrArgs a{N, HC, LR, (uint32_t) T, 1u, eps_}; + D(P(lib_, "gr_norm", {}, "gr_norm"), {mtl::Arg::B(p.R), mtl::Arg::B(p.bo), mtl::Arg::B(p.inj[ci]), mtl::Arg::B(W_, tn.offset), + mtl::Arg::B(p.YE), mtl::Arg::V(a)}, HC, T, 256); + } + const double t0 = now_ms(); + if (stream_) { + seg_T_ = T; + cut_ids_ = p.ids.as(); + cx_->run_graph_seg(ds, cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this); + cut_ids_ = nullptr; + } else { + cx_->run_graph(ds); + } + pg_ms_ += now_ms() - t0; + pg_tok_ += T; + static const bool trace = std::getenv("STRATA_METAL_TRACE") != nullptr; + if (trace) std::fprintf(stderr, "strata metal: prompt chunk pos %lld, %d rows: %.0f ms\n", (long long) pos0, T, now_ms() - t0); + // the window's state after the chunk: the conv history turned, the PLE history and token window advanced + hp_ = 1 - hp_; + if (ple_on_) { + using namespace strata::kernels; + float* hist = pleHist_.as(); + const float* nw = p.normW.as(); + for (int c = 0; c < NG_HC_DIM; ++c) { + float col[NG_HIST]; + for (int r = 0; r < NG_HIST; ++r) { + const int src = T - NG_HIST + r; // row r of the new history: chunk row src, or the old history + col[r] = src >= 0 ? nw[(size_t) src * NG_HC_DIM + c] : hist[(size_t) c * NG_HIST + (size_t) (r + T)]; + } + std::memcpy(hist + (size_t) c * NG_HIST, col, sizeof col); + } + } + if (T >= 2) { ss.ple_prev[0] = tokens[T - 2]; ss.ple_prev[1] = tokens[T - 1]; } + else { ss.ple_prev[0] = ss.ple_prev[1]; ss.ple_prev[1] = tokens[0]; } + mtp_src_ = &p.R; // the draft layer's prompt K/V read this chunk's final rows + return true; +} + +// ---------------------------------------------------------------------------------------------- the hook's entry points +namespace { +std::string g_shard1; +} +void window_register(void* w); +bool mtp_load_pending(Window* w, std::string& err); +void set_model_shard(const char* shard1) { g_shard1 = shard1 ? shard1 : ""; } +void set_batch_slots(int n) { g_slot_cap = n > 0 ? std::min(n, 8) : 0; } + +// One decode window per process: the prompt path (Prefill::init) may bring it up before the Verifier does. +namespace { +Window* g_the_window = nullptr; +} +void* window_create() { + if (g_the_window == nullptr) g_the_window = new Window(); + return g_the_window; +} +void window_destroy(void* w) { + if (w != nullptr && w == g_the_window) { + if (g_stream_window == g_the_window) { + g_the_window->stream_report(stderr); + g_stream_window = nullptr; + } + delete g_the_window; + g_the_window = nullptr; + window_register(nullptr); + } +} +bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::SessionState& ss, int max_t, int64_t max_cells, + int64_t* n_vocab, std::string& err) { + Window* x = (Window*) w; + if (!x->ready()) { + if (g_shard1.empty()) { err = "metal window: the model shard is not set (--native SHARD1)"; return false; } + if (!x->init(g, ss, g_shard1, MAXT, max_cells, err)) return false; + window_register(w); + if (!mtp_load_pending(x, err)) return false; + } + x->set_max_t(max_t); + *n_vocab = x->vocab(); + return true; +} +bool window_ensure(const strata::core::ModelGeometry& g, strata::core::SessionState& ss, int64_t max_cells, std::string& err) { + int64_t nv = 0; + return window_init(window_create(), g, ss, MAXT, max_cells, &nv, err); +} +bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, std::string& err) { + // STRATA_LOGPOS (the quant yardstick, the KL harness): every window's head logits are read, also the prompt's + // windows that sample nothing (head_sampling off) + static const bool logpos = std::getenv("STRATA_LOGPOS") != nullptr; + return ((Window*) w)->run(T, tokens, pos0, head || logpos, out, err); +} +void window_reset(const strata::core::SessionState* ss) { + if (g_the_window == nullptr || !g_the_window->ready()) return; + g_the_window->activate(g_the_window->set_of(ss)); + g_the_window->reset(); + g_the_window->activate(0); +} +bool window_checkpoint_save(strata::core::ConversationCheckpoint& c, const strata::core::SessionState* ss, std::string& err) { + if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } + g_the_window->activate(g_the_window->set_of(ss)); // a batch slot's session: that slot's running state + const bool ok = g_the_window->checkpoint_save(c, err); + g_the_window->activate(0); + return ok; +} +bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, const strata::core::SessionState* ss, std::string& err) { + if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } + g_the_window->activate(g_the_window->set_of(ss)); + const bool ok = g_the_window->checkpoint_restore(c, err); + g_the_window->activate(0); + return ok; +} +bool window_init_slots(void* w, const std::vector& slots, std::string& err) { + return ((Window*) w)->slots_init(slots, err); +} +bool window_run_slots(void* w, const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err) { + const double t0 = now_ms(); + const bool ok = ((Window*) w)->run_slots(rows, S, tokens, pos, out, err); + ((Window*) w)->prof_add(true, now_ms() - t0); + return ok; +} +bool window_slot_copy_main(bool into_slot, int slot, int64_t upto, bool kv_only, std::string& err) { + if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } + return g_the_window->slot_copy(into_slot, slot, upto, kv_only, err); +} +bool window_commit(void* w, int n_keep, std::string& err) { + const double t0 = now_ms(); + const bool ok = ((Window*) w)->commit(n_keep, err); + ((Window*) w)->prof_add(false, now_ms() - t0); + return ok; +} +bool window_copy_logits(void* w, int t, float* host) { return ((Window*) w)->copy_logits(t, host); } + +// the MTP hook (src/core/mtp.cpp under STRATA_GPU_METAL): the drafter lives beside the decode window it reads +namespace { +Window* g_window = nullptr; +} +void window_register(void* w) { g_window = (Window*) w; } +// the engine loads the drafter before the decode window exists: the path is kept and the window loads it at its init +std::string g_mtp_path; +int g_mtp_max_t = 0; +bool mtp_load(const std::string& rt_dir, int max_t, std::string& err) { + (void) err; + // the MTP model in GGUF form (tools/mtp_pack.py) beside the rt directory: STRATA_METAL_MTP_GGUF, else + // /../mtp-q2_0.gguf + std::string path; + if (const char* e = std::getenv("STRATA_METAL_MTP_GGUF")) path = e; + else { + std::string d = rt_dir; + while (!d.empty() && d.back() == '/') d.pop_back(); + const size_t s = d.rfind('/'); + path = (s == std::string::npos ? std::string(".") : d.substr(0, s)) + "/mtp-q2_0.gguf"; + } + g_mtp_path = path; + g_mtp_max_t = max_t; + return true; +} +bool mtp_load_pending(Window* w, std::string& err) { + return g_mtp_path.empty() || w->mtp_load(g_mtp_path, g_mtp_max_t, err); +} +// Prefill::run's Metal side: the prompt read through the window, MAXT rows at a time (no head), each window +// committed; on_chunk (the MTP's K/V) after each, with the window's final residual rows in place. +bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, + const std::function& on_chunk, std::string& err); +bool prompt_run(const int64_t* tokens, int64_t n, int64_t pos0, const float* const* embd_rows, + const std::function& on_chunk, std::string& err) { + struct Rows { + explicit Rows(const float* const* r) { g_embd_rows = r; } + ~Rows() { g_embd_rows = nullptr; } + } rows(embd_rows); + return prompt_run_rows(tokens, n, pos0, on_chunk, err); +} +bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, + const std::function& on_chunk, std::string& err) { + if (g_window == nullptr) { err = "metal prompt: the decode window is not up yet"; return false; } + static const bool force_windows = std::getenv("STRATA_METAL_PROMPT") && std::string(std::getenv("STRATA_METAL_PROMPT")) == "window"; + static const char* latent_prompt = std::getenv("STRATA_METAL_LATENT_PROMPT"); + if (g_window->latent() && !force_windows && latent_prompt && std::strcmp(latent_prompt, "1") == 0 && n >= 128) + return g_window->prompt(tokens, n, pos0, on_chunk, err); + const bool windows = force_windows || g_window->latent(); + static const int64_t min_rows = std::getenv("STRATA_METAL_PROMPT_MIN") ? std::atoll(std::getenv("STRATA_METAL_PROMPT_MIN")) : 32; + if (!windows && !g_window->streaming() && n >= min_rows) return g_window->prompt(tokens, n, pos0, on_chunk, err); + if (!windows && g_window->streaming()) { + const char* e = std::getenv("STRATA_METAL_PROMPT_GEMM"); + const char* em = std::getenv("STRATA_METAL_PROMPT_GEMM_MIN"); + const int64_t mn = em ? std::atoll(em) : 64; + if (!(e && std::string(e) == "0") && n >= mn) { + if (!g_window->prompt_gemm_init(err)) return false; + const int C = g_window->prompt_chunk_rows(); + std::vector tk32((size_t) C); + g_window->prompt_region(true); + bool ok = true; + for (int64_t c = 0; c < n && ok; c += C) { + const int k = (int) std::min(C, n - c); + for (int i = 0; i < k; ++i) tk32[(size_t) i] = (int32_t) tokens[c + i]; + ok = g_window->prompt_chunk(tk32.data(), k, pos0 + c, err) && (!on_chunk || on_chunk(nullptr, k, pos0 + c, err)); + } + g_window->prompt_region(false); + return ok; + } + } + int32_t tk[MAXT], out[MAXT]; + for (int64_t c = 0; c < n; c += MAXT) { + const int k = (int) std::min(MAXT, n - c); + for (int i = 0; i < k; ++i) tk[i] = (int32_t) tokens[c + i]; + if (!g_window->run(k, tk, pos0 + c, false, out, err) || !g_window->commit(k, err)) return false; + if (on_chunk && !on_chunk(nullptr, k, pos0 + c, err)) return false; + } + if (n >= 1024) g_window->report_profile(); + return true; +} +bool mtp_prefill(const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err) { + if (g_window == nullptr) { err = "metal mtp: no decode window"; return false; } + return g_window->mtp_prefill(next_tokens, n, cell0, err); +} +bool mtp_draft(int T, const int32_t* tokens, int64_t p, int a, int max_steps, int32_t* drafts, float* probs, float min_p, + int* n_drafts, std::string& err) { + if (g_window == nullptr) { err = "metal mtp: no decode window"; return false; } + const double t0 = now_ms(); + const bool ok = g_window->mtp_draft(T, tokens, p, a, max_steps, drafts, probs, min_p, n_drafts, err); + g_window->prof_add(true, now_ms() - t0); + return ok; +} + +#include "prompt_chain.inc" + +} // namespace strata::metal_engine +namespace strata::gpu { void* host_view(const void* device_pointer); } +namespace strata::kernels { +void mrope_table_set(const int32_t* device_table) { + strata::metal_engine::g_mrope_dev = device_table; + strata::metal_engine::g_mrope_host = + device_table ? (const int32_t*) strata::gpu::host_view(device_table) : nullptr; +} +const int32_t* mrope_table() { return strata::metal_engine::g_mrope_dev; } +} // namespace strata::kernels +namespace strata::metal_engine { + +} // namespace strata::metal_engine diff --git a/edge0-strata/metal/src/gguf_lite.cpp b/edge0-strata/metal/src/gguf_lite.cpp new file mode 100644 index 00000000..44937fba --- /dev/null +++ b/edge0-strata/metal/src/gguf_lite.cpp @@ -0,0 +1,117 @@ +// metal/src/gguf_lite.cpp - see gguf_lite.h. +#include "gguf_lite.h" + +#include "ref_quant.h" + +#include +#include +#include +#include + +#include +#include + +namespace gguf { + +namespace { +struct Reader { + const uint8_t* p; + const uint8_t* end; + bool ok = true; + template T get() { + T v{}; + if (p + sizeof(T) > end) { ok = false; return v; } + std::memcpy(&v, p, sizeof(T)); + p += sizeof(T); + return v; + } + std::string str() { + const uint64_t n = get(); + if (!ok || p + n > end) { ok = false; return {}; } + std::string s((const char*)p, n); + p += n; + return s; + } +}; +// value types: 0 u8 1 i8 2 u16 3 i16 4 u32 5 i32 6 f32 7 bool 8 string 9 array 10 u64 11 i64 12 f64 +size_t scalar_size(uint32_t t) { + switch (t) { + 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; + } + return 0; +} +double scalar_num(Reader& r, uint32_t t) { + switch (t) { + case 0: return r.get(); case 1: return r.get(); case 2: return r.get(); + case 3: return r.get(); case 4: return r.get(); case 5: return r.get(); + case 6: return r.get(); case 7: return r.get(); case 10: return (double)r.get(); + case 11: return (double)r.get(); case 12: return r.get(); + } + r.ok = false; + return 0; +} +} // namespace + +bool File::open(const std::string& p, std::string& err) { + path = p; + const int fd = ::open(p.c_str(), O_RDONLY); + if (fd < 0) { err = "cannot open " + p; return false; } + struct stat st; + fstat(fd, &st); + file_bytes = (size_t)st.st_size; + const size_t page = (size_t)getpagesize(); + map_bytes = (file_bytes + page - 1) / page * page; + map = mmap(nullptr, map_bytes, PROT_READ, MAP_SHARED, fd, 0); + ::close(fd); + if (map == MAP_FAILED) { map = nullptr; err = "mmap failed: " + p; return false; } + Reader r{(const uint8_t*)map, (const uint8_t*)map + file_bytes}; + if (r.get() != 0x46554747u) { err = "not a GGUF file"; return false; } + if (r.get() != 3) { err = "GGUF version is not 3"; return false; } + const uint64_t n_t = r.get(), n_kv = r.get(); + uint64_t align = 32; + for (uint64_t i = 0; i < n_kv && r.ok; ++i) { + const std::string key = r.str(); + const uint32_t t = r.get(); + if (t == 8) meta_str[key] = r.str(); + else if (t == 9) { + const uint32_t et = r.get(); + const uint64_t n = r.get(); + if (et == 8) for (uint64_t k = 0; k < n && r.ok; ++k) r.str(); + else r.p += n * scalar_size(et); + } else { + const double v = scalar_num(r, t); + meta_num[key] = v; + if (key == "general.alignment") align = (uint64_t)v; + } + } + std::vector ts; + for (uint64_t i = 0; i < n_t && r.ok; ++i) { + Tensor t; + t.name = r.str(); + const uint32_t nd = r.get(); + for (uint32_t k = 0; k < nd; ++k) t.shape.push_back(r.get()); + t.type = (int)r.get(); + t.offset = r.get(); + ts.push_back(t); + } + if (!r.ok) { err = "truncated GGUF header"; return false; } + const uint64_t pos = (uint64_t)(r.p - (const uint8_t*)map); + const uint64_t data = (pos + align - 1) / align * align; + for (Tensor& t : ts) { + t.offset += data; + uint64_t n = 1; + for (uint64_t d : t.shape) n *= d; + t.bytes = n / ref::block_elems(t.type) * ref::block_bytes(t.type); + tensors[t.name] = t; + } + return true; +} + +File::~File() { + if (map) munmap(map, map_bytes); +} + +} // namespace gguf diff --git a/edge0-strata/metal/src/gguf_lite.h b/edge0-strata/metal/src/gguf_lite.h new file mode 100644 index 00000000..0a2b5fbb --- /dev/null +++ b/edge0-strata/metal/src/gguf_lite.h @@ -0,0 +1,37 @@ +// metal/src/gguf_lite.h - a minimal GGUF v3 reader for the Metal test programs: tensor names, shapes, types and +// file offsets, the file mmapped read-only (no copy; wrapped as one MTLBuffer by the caller). +#pragma once + +#include +#include +#include +#include +#include + +namespace gguf { + +struct Tensor { + std::string name; + std::vector shape; // ggml order: shape[0] is the contiguous (row) dimension + int type = -1; + uint64_t offset = 0; // from the start of the FILE (data start + tensor offset) + uint64_t bytes = 0; +}; + +struct File { + std::string path; + void* map = nullptr; // page-aligned mapping of the whole file + size_t map_bytes = 0; // rounded up to the page size + size_t file_bytes = 0; + std::map tensors; + std::map meta_str; + std::map meta_num; + bool open(const std::string& path, std::string& err); + const Tensor* find(const std::string& name) const { + auto it = tensors.find(name); + return it == tensors.end() ? nullptr : &it->second; + } + ~File(); +}; + +} // namespace gguf diff --git a/edge0-strata/metal/src/gpu/runtime.cpp b/edge0-strata/metal/src/gpu/runtime.cpp new file mode 100644 index 00000000..43e2d30e --- /dev/null +++ b/edge0-strata/metal/src/gpu/runtime.cpp @@ -0,0 +1,857 @@ +// metal/src/gpu/runtime.cpp - include/strata/gpu/runtime.hpp on Metal (metal-cpp). Skeleton of MAC-PLAN Phase 1. +// +// Memory model (measured on the M5 Max, phaseA-tests/mac-design.md): +// - every allocation is one MTLBuffer, storageModeShared, hazard tracking untracked (the engine orders its own work +// on streams, as on CUDA); the "device pointer" handed to the host code is the buffer's gpuAddress, which differs +// from contents() on Apple GPUs, so host code that wrongly dereferences a device pointer faults instead of +// silently working. host_view() is the sanctioned CPU view. +// - all buffers sit in one MTLResidencySet attached to every queue, so kernels may take raw gpuAddress pointers. +// Streams: a command queue + the open command buffer + the open compute encoder + a shared-event timeline. Work is +// committed at event records, synchronizations, host functions and every kFlushEvery dispatches. +// Graphs: begin_capture switches the stream to recording; dispatch/copy/memset/event ops are kept as data and +// re-encoded by graph_launch (encoding measured at ~0.17 us per dispatch on the M5 Max). +#define NS_PRIVATE_IMPLEMENTATION +#define MTL_PRIVATE_IMPLEMENTATION +#include +#include +#include + +#include "strata/gpu/metal.hpp" +#include "strata/gpu/runtime.hpp" + +#include + +#include +#include +#include +#if defined(STRATA_DS4_RUNTIME) +#include +#endif +#include +#include +#include +#if defined(STRATA_DS4_RUNTIME) +#include // recursive_mutex +#else +#include +#endif +#include +#if defined(STRATA_DS4_RUNTIME) +#include +#endif +#include + +namespace strata::gpu { + +namespace { + +constexpr int kFlushEvery = 4096; // dispatches per command buffer before an automatic commit + +struct Alloc { + MTL::Buffer* buf = nullptr; + uint64_t gpu = 0; + uint8_t* cpu = nullptr; + size_t bytes = 0; + bool host = false; // from host_alloc (the host code holds the CPU address) + bool purged = false; // purge_device_allocations: its memory given back (any later copy is an error) +}; + +struct Op { + enum Kind { Dispatch, Copy, Fill, Record, Wait } kind = Dispatch; + MTL::ComputePipelineState* pipe = nullptr; + uint8_t args[256]; + size_t arg_bytes = 0; + std::vector big_args; + metal::Grid groups, threads; + size_t tg_bytes = 0; + void* dst = nullptr; + const void* src = nullptr; + size_t bytes = 0; + int value = 0; + Event ev = nullptr; +}; + +} // namespace + +struct EventImpl { + MTL::SharedEvent* ev = nullptr; // the recording stream's timeline + uint64_t value = 0; // 0: never recorded (complete, as a fresh CUDA event) +#if defined(STRATA_DS4_RUNTIME) + void* token = nullptr; // a mark on the context's queue (ds4_ctx_mark), released by the wait +#endif +}; + +struct GraphImpl { + std::vector ops; +}; +struct GraphExecImpl { + std::vector ops; +}; + +struct StreamImpl { + MTL::CommandQueue* q = nullptr; + MTL::CommandBuffer* cb = nullptr; + MTL::ComputeCommandEncoder* enc = nullptr; + MTL::SharedEvent* tl = nullptr; + uint64_t submitted = 0; // the timeline value the last committed command buffer signals + int pending = 0; // dispatches in the open command buffer + GraphImpl* rec = nullptr; +}; + +namespace { + +struct Runtime { + MTL::Device* dev = nullptr; + MTL::ResidencySet* rs = nullptr; + std::mutex mu; + std::map by_gpu; // gpuAddress -> allocation + std::map by_cpu; // contents() -> gpuAddress + StreamImpl* legacy = nullptr; + std::map pipes; + std::map libs; + metal::Stats stats; +}; + +Runtime& rt() { + static Runtime* r = [] { + auto* x = new Runtime; + x->dev = MTL::CreateSystemDefaultDevice(); + NS::Error* e = nullptr; + MTL::ResidencySetDescriptor* d = MTL::ResidencySetDescriptor::alloc()->init(); + x->rs = x->dev->newResidencySet(d, &e); + d->release(); + return x; + }(); + return *r; +} + +thread_local Error t_last = kSuccess; +Error fail(Error e) { + if (e != kSuccess && e != kNotReady) t_last = e; + return e; +} + +#if defined(STRATA_DS4_RUNTIME) +static std::vector g_all_streams; // every stream's queue holds pending copies until it is committed +static std::mutex g_all_streams_mu; + +#endif +StreamImpl* make_stream() { + auto* s = new StreamImpl; +#if defined(STRATA_DS4_RUNTIME) + { std::lock_guard lk(g_all_streams_mu); g_all_streams.push_back(s); } +#endif + s->q = rt().dev->newCommandQueue(); + s->q->addResidencySet(rt().rs); + s->tl = rt().dev->newSharedEvent(); + s->tl->setSignaledValue(0); + return s; +} + +StreamImpl* S(Stream s) { + if (s) return s; + Runtime& r = rt(); + std::lock_guard lk(r.mu); + if (!r.legacy) r.legacy = make_stream(); + return r.legacy; +} + +/// The allocation that holds GPU address `p` (or CPU address `p` of a host allocation), with the offset into it. +const Alloc* find(const void* p, size_t* off) { + Runtime& r = rt(); + std::lock_guard lk(r.mu); + const uint64_t a = (uint64_t) (uintptr_t) p; + auto it = r.by_gpu.upper_bound(a); + if (it != r.by_gpu.begin()) { + --it; + if (a < it->first + it->second.bytes) { *off = a - it->first; return &it->second; } + } + auto ic = r.by_cpu.upper_bound((uintptr_t) p); + if (ic != r.by_cpu.begin()) { + --ic; + const Alloc& al = r.by_gpu[ic->second]; + if ((uintptr_t) p < ic->first + al.bytes) { *off = (uintptr_t) p - ic->first; return &al; } + } + return nullptr; +} + +#if defined(STRATA_DS4_RUNTIME) +// The engine's host threads (the prompt path, the loader, the expert cache) have no autorelease pool, and metal-cpp's +// commandBuffer() / computeCommandEncoder() / blitCommandEncoder() hand back autoreleased objects: with no pool that last reference was +// never given back, so every command buffer and encoder stayed alive with the buffers it had named (~800 command buffers a +// long prompt). A pool around each creation (and, for the blits, around the whole encode) gives it back. +struct ScopedPool { + NS::AutoreleasePool* p; + ScopedPool() : p(NS::AutoreleasePool::alloc()->init()) {} + ~ScopedPool() { p->release(); } +}; +#endif +MTL::CommandBuffer* open_cb(StreamImpl* s) { + if (!s->cb) { +#if defined(STRATA_DS4_RUNTIME) + ScopedPool pool; +#endif + s->cb = s->q->commandBufferWithUnretainedReferences(); + s->cb->retain(); + ++rt().stats.command_buffers; + } + return s->cb; +} +#if defined(STRATA_DS4_RUNTIME) +// StreamImpl's encoder/command-buffer state is touched from the loader, the prefetcher and the compute path. +// Without this lock two encoders can be open on one command buffer at once, which Metal asserts on. +static std::recursive_mutex g_stream_mu; // encode_copy/flush nest + +#endif +void end_encoder(StreamImpl* s) { +#if defined(STRATA_DS4_RUNTIME) + if (s->enc) { if (std::getenv("STRATA_GPU_TRACE")) std::fprintf(stderr, "[gpu] close compute enc %p\n", (void*) s->enc); + s->enc->endEncoding(); s->enc->release(); s->enc = nullptr; } +#else + if (s->enc) { s->enc->endEncoding(); s->enc = nullptr; } +#endif +} +MTL::ComputeCommandEncoder* encoder(StreamImpl* s) { +#if defined(STRATA_DS4_RUNTIME) + std::lock_guard lk(g_stream_mu); + if (!s->enc) { + MTL::CommandBuffer* cb = open_cb(s); + { ScopedPool pool; s->enc = cb->computeCommandEncoder(); s->enc->retain(); } // (released in end_encoder) + if (std::getenv("STRATA_GPU_TRACE")) std::fprintf(stderr, "[gpu] open compute enc %p (cb %p)\n", (void*) s->enc, (void*) s->cb); } // serial dispatch: each dispatch sees the last +#else + if (!s->enc) s->enc = open_cb(s)->computeCommandEncoder(); // serial dispatch: each dispatch sees the last +#endif + return s->enc; +} +void flush(StreamImpl* s) { +#if defined(STRATA_DS4_RUNTIME) + std::lock_guard lk(g_stream_mu); +#endif + if (!s->cb) return; + end_encoder(s); +#if defined(STRATA_DS4_RUNTIME) + if (std::getenv("STRATA_GPU_TRACE")) std::fprintf(stderr, "[gpu] signal %llu on cb %p (enc %p)\n", + (unsigned long long) (s->submitted + 1), (void*) s->cb, (void*) s->enc); +#endif + s->cb->encodeSignalEvent(s->tl, ++s->submitted); + s->cb->commit(); + s->cb->release(); + s->cb = nullptr; + s->pending = 0; +} +#if defined(STRATA_DS4_RUNTIME) +extern "C" void ds4_metal_flush(); +extern "C" bool ds4_ctx_copy(void* dst, const void* src, size_t n) __attribute__((weak_import)); +extern "C" bool ds4_ctx_fill(void* dst, int value, size_t n) __attribute__((weak_import)); +extern "C" void* ds4_ctx_mark() __attribute__((weak_import)); +extern "C" void ds4_ctx_wait(void* token) __attribute__((weak_import)); +extern "C" void ds4_ctx_drop(void* token) __attribute__((weak_import)); +static bool ctx_queue_ops() { // DS4_STREAM_COPIES=1: the old way (a copy on a stream's own queue, a full barrier each) + static const bool on = ds4_ctx_copy != nullptr && std::getenv("DS4_STREAM_COPIES") == nullptr; + return on; +} + +#endif +Error wait_value(MTL::SharedEvent* ev, uint64_t v) { + while (ev->signaledValue() < v) std::this_thread::yield(); + ++rt().stats.host_syncs; + return kSuccess; +} + +void check_purged(const Alloc* a, const char* what) { + if (a && a->purged) { + std::fprintf(stderr, "strata metal runtime: a %s touches a device allocation of %.1f MB that was given back " + "(STRATA_METAL_KEEP_HOST_WEIGHTS=1 keeps it)\n", what, a->bytes / 1e6); + void* fr[48]; + backtrace_symbols_fd(fr, backtrace(fr, 48), 2); // which host code still reads it + std::abort(); + } +} + +Error encode_copy(StreamImpl* s, void* dst, const void* src, size_t n) { +#if defined(STRATA_DS4_RUNTIME) + std::lock_guard lk(g_stream_mu); +#endif + size_t od = 0, os = 0; + const Alloc* ad = find(dst, &od); + const Alloc* as = find(src, &os); + check_purged(ad, "copy"); + check_purged(as, "copy"); + if (ad && as) { // buffer to buffer: a blit, in stream order + end_encoder(s); +#if defined(STRATA_DS4_RUNTIME) + ScopedPool pool; +#endif + MTL::BlitCommandEncoder* b = open_cb(s)->blitCommandEncoder(); +#if defined(STRATA_DS4_RUNTIME) + if (std::getenv("STRATA_GPU_TRACE")) std::fprintf(stderr, "[gpu] blit copy %zu\n", n); +#endif + b->copyFromBuffer(as->buf, os, ad->buf, od, n); + b->endEncoding(); + ++rt().stats.blits; + return kSuccess; + } + // ordinary host memory on one side: the stream's earlier work must be complete first (the conservative CUDA + // pageable-memory semantics: the call returns when the copy is done) + flush(s); + wait_value(s->tl, s->submitted); + uint8_t* d = ad ? ad->cpu + od : (uint8_t*) dst; + const uint8_t* sp = as ? as->cpu + os : (const uint8_t*) src; + std::memcpy(d, sp, n); + return kSuccess; +} + +Error encode_fill(StreamImpl* s, void* p, int value, size_t n) { +#if defined(STRATA_DS4_RUNTIME) + std::lock_guard lk(g_stream_mu); +#endif + size_t off = 0; + const Alloc* a = find(p, &off); + if (!a) return fail(kErrorInvalidValue); + check_purged(a, "fill"); + end_encoder(s); +#if defined(STRATA_DS4_RUNTIME) + ScopedPool pool; +#endif + MTL::BlitCommandEncoder* b = open_cb(s)->blitCommandEncoder(); + b->fillBuffer(a->buf, NS::Range::Make(off, n), (uint8_t) value); + b->endEncoding(); + ++rt().stats.blits; + return kSuccess; +} + +Error encode_dispatch(StreamImpl* s, const Op& o) { +#if defined(STRATA_DS4_RUNTIME) + std::lock_guard lk(g_stream_mu); +#endif + MTL::ComputeCommandEncoder* e = encoder(s); + e->setComputePipelineState(o.pipe); + if (o.big_args.empty()) e->setBytes(o.args, o.arg_bytes, 0); + else e->setBytes(o.big_args.data(), o.big_args.size(), 0); + if (o.tg_bytes) e->setThreadgroupMemoryLength(o.tg_bytes, 0); + e->dispatchThreadgroups(MTL::Size(o.groups.x, o.groups.y, o.groups.z), MTL::Size(o.threads.x, o.threads.y, o.threads.z)); + ++rt().stats.dispatches; + if (++s->pending >= kFlushEvery) flush(s); + return kSuccess; +} + +Error record_event(EventImpl* e, StreamImpl* s) { + flush(s); // CUDA submits on launch: a recorded event must become complete without a further call + e->ev = s->tl; + e->value = s->submitted; + return kSuccess; +} +Error wait_event(StreamImpl* s, EventImpl* e) { + if (!e || e->value == 0) return kSuccess; + end_encoder(s); + open_cb(s)->encodeWait(e->ev, e->value); + return kSuccess; +} + +Error run_op(StreamImpl* s, const Op& o) { + switch (o.kind) { + case Op::Dispatch: return encode_dispatch(s, o); + case Op::Copy: return encode_copy(s, o.dst, o.src, o.bytes); + case Op::Fill: return encode_fill(s, o.dst, o.value, o.bytes); + case Op::Record: return record_event(o.ev, s); + case Op::Wait: return wait_event(s, o.ev); + } + return kSuccess; +} + +// STRATA_METAL_ALLOC_LOG=1: every allocation of 16 MB or more, with the running total (memory budgets on small Macs) +void alloc_log(const char* what, size_t n) { + static const bool on = std::getenv("STRATA_METAL_ALLOC_LOG") != nullptr; + static std::atomic tot{0}; + const uint64_t t = tot += n; + if (on && n >= (16u << 20)) std::fprintf(stderr, "strata metal runtime: %s %.1f MB (total %.2f GB)\n", what, n / 1e6, t / 1e9); +} + +Error new_alloc(void** p, size_t bytes, bool host) { + Runtime& r = rt(); +#if defined(STRATA_DS4_RUNTIME) + const size_t pg = 16384, want = bytes ? bytes : 1; + const size_t n = std::getenv("DS4_NO_PAD") ? want : ((want + pg - 1) / pg) * pg; +#else + const size_t n = bytes ? bytes : 1; +#endif + MTL::Buffer* b = r.dev->newBuffer(n, MTL::ResourceStorageModeShared | MTL::ResourceHazardTrackingModeUntracked); + if (!b) return fail(kErrorMemoryAllocation); + Alloc a{b, b->gpuAddress(), (uint8_t*) b->contents(), n, host}; + alloc_log(host ? "host alloc" : "device alloc", n); + { + std::lock_guard lk(r.mu); + r.by_gpu[a.gpu] = a; + r.by_cpu[(uintptr_t) a.cpu] = a.gpu; + r.rs->addAllocation(b); + r.rs->commit(); + } + *p = host ? (void*) a.cpu : (void*) (uintptr_t) a.gpu; + return kSuccess; +} + +Error free_alloc(void* p) { + if (!p) return kSuccess; + Runtime& r = rt(); + size_t off = 0; + const Alloc* a = find(p, &off); + if (!a || off != 0) return fail(kErrorInvalidValue); +#if defined(STRATA_DS4_RUNTIME) + if (std::getenv("DS4_TRACE_FREE")) std::fprintf(stderr, "gpu free %zu bytes (host=%d) at %p\n", a->bytes, (int) a->host, p); +#endif + device_synchronize(); // as cudaFree: the GPU is done with it + std::lock_guard lk(r.mu); + MTL::Buffer* b = a->buf; + const uint64_t g = a->gpu; + r.by_cpu.erase((uintptr_t) a->cpu); + r.by_gpu.erase(g); + r.rs->removeAllocation(b); + r.rs->commit(); + b->release(); + return kSuccess; +} + +} // namespace + +// ---------------------------------------------------------------------------------------------------- runtime.hpp +#if defined(STRATA_DS4_RUNTIME) +void* buffer_for_cpu_impl(const void* cpu, size_t bytes, size_t* off) { + Runtime& r = rt(); + std::lock_guard lk(r.mu); + auto ic = r.by_cpu.upper_bound((uintptr_t) cpu); + if (ic == r.by_cpu.begin()) return nullptr; + --ic; + const Alloc& a = r.by_gpu[ic->second]; + if ((uintptr_t) cpu < ic->first || (uintptr_t) cpu + bytes > ic->first + a.bytes) return nullptr; + *off = (uintptr_t) cpu - ic->first; + return (void*) a.buf; +} + +MTL::ResidencySet* residency_set_impl() { return rt().rs; } + +#endif +void* host_view(const void* p) { + size_t off = 0; + const Alloc* a = find(p, &off); + return a ? a->cpu + off : nullptr; +} + +// The Metal decode window reads the model itself; the device copies of the weights the shared host code uploads for +// the CUDA-shaped kernels (the pack's dense pool, the native head, ...) are never read on Metal. On a small Mac their +// memory is worth more as expert arena: every device (not host) allocation of min_bytes or more is made purgeable-empty +// and dropped from the residency set. Returns the bytes given back. +uint64_t purge_device_allocations(size_t min_bytes) { + Runtime& r = rt(); + device_synchronize(); + std::lock_guard lk(r.mu); + uint64_t n = 0; + for (auto& [gpu, a] : r.by_gpu) { + if (a.host || a.purged || a.bytes < min_bytes) continue; + r.rs->removeAllocation(a.buf); + a.buf->setPurgeableState(MTL::PurgeableStateEmpty); + a.purged = true; + n += a.bytes; + } + r.rs->commit(); + return n; +} + +Error get_last_error() { const Error e = t_last; t_last = kSuccess; return e; } +Error peek_last_error() { return t_last; } +const char* error_string(Error e) { + switch (e) { + case kSuccess: return "no error"; + case kNotReady: return "device not ready"; + case kErrorInvalidValue: return "invalid argument"; + case kErrorMemoryAllocation: return "out of memory"; + case kErrorUnsupported: return "operation not supported on Metal"; + case kErrorLaunchFailure: return "command buffer failed"; + } + return "unknown Metal runtime error"; +} + +Error get_device(int* dev) { *dev = 0; return kSuccess; } +Error set_device(int dev) { return dev == 0 ? kSuccess : fail(kErrorInvalidValue); } +Error device_count(int* n) { *n = rt().dev ? 1 : 0; return kSuccess; } +Error device_synchronize() { +#if defined(STRATA_DS4_RUNTIME) + ds4_metal_flush(); +#endif + Runtime& r = rt(); + StreamImpl* s = S(nullptr); + flush(s); + return wait_value(s->tl, s->submitted); + (void) r; +} +Error mem_get_info(size_t* free_b, size_t* total_b) { + // the GPU's budget is the working set the OS grants this process; what is free is that minus what is allocated + MTL::Device* d = rt().dev; + const uint64_t budget = d->recommendedMaxWorkingSetSize(); + const uint64_t used = d->currentAllocatedSize(); + *total_b = (size_t) budget; + *free_b = (size_t) (used < budget ? budget - used : 0); + return kSuccess; +} +Error device_integrated(int* v, int) { *v = 1; return kSuccess; } +Error device_sm_count(int* v, int) { *v = 0; return fail(kErrorUnsupported); } // Apple reports no core count +Error device_clock_khz(int* v, int) { *v = 0; return fail(kErrorUnsupported); } +Error runtime_version(int* v) { *v = 0; return kSuccess; } +Error device_can_access_peer(int* v, int, int) { *v = 0; return kSuccess; } +Error device_enable_peer_access(int, unsigned) { return fail(kErrorUnsupported); } +Error get_device_properties(DeviceProp* p, int dev) { + if (dev != 0) return fail(kErrorInvalidValue); + *p = DeviceProp{}; + MTL::Device* d = rt().dev; + std::snprintf(p->name, sizeof p->name, "%s", d->name()->utf8String()); + std::snprintf(p->gcnArchName, sizeof p->gcnArchName, "%s", d->supportsFamily(MTL::GPUFamilyApple10) ? "apple10" : "apple9-or-older"); + p->warpSize = 32; // the SIMD-group width of every Apple GPU + p->totalGlobalMem = (size_t) d->recommendedMaxWorkingSetSize(); + return kSuccess; +} + +Error malloc_device_raw(void** p, size_t bytes) { return new_alloc(p, bytes, false); } +Error free_device(void* p) { return free_alloc(p); } +Error host_alloc_raw(void** p, size_t bytes, unsigned) { return new_alloc(p, bytes, true); } +Error free_host(void* p) { + { + Runtime& r = rt(); + std::lock_guard lk(r.mu); + auto it = r.by_cpu.find((uintptr_t) p); + if (it == r.by_cpu.end()) return fail(kErrorInvalidValue); + } + return free_alloc(p); +} +Error host_register(void* p, size_t bytes, unsigned) { + // existing host memory (an mmapped file, an arena) as a NoCopy buffer: page-aligned start and length required + const uintptr_t pg = 16384; + if (((uintptr_t) p % pg) != 0) return fail(kErrorInvalidValue); + Runtime& r = rt(); + const size_t n = (bytes + pg - 1) / pg * pg; + MTL::Buffer* b = r.dev->newBuffer(p, n, MTL::ResourceStorageModeShared | MTL::ResourceHazardTrackingModeUntracked, nullptr); + if (!b) return fail(kErrorMemoryAllocation); + Alloc a{b, b->gpuAddress(), (uint8_t*) b->contents(), n, true}; + alloc_log("host register", n); + std::lock_guard lk(r.mu); + r.by_gpu[a.gpu] = a; + r.by_cpu[(uintptr_t) a.cpu] = a.gpu; + r.rs->addAllocation(b); + r.rs->commit(); + return kSuccess; +} +Error host_unregister(void* p) { return free_alloc(p); } +Error host_device_pointer_raw(void** dev, void* host) { + size_t off = 0; + const Alloc* a = find(host, &off); + if (!a) return fail(kErrorInvalidValue); + *dev = (void*) (uintptr_t) (a->gpu + off); + return kSuccess; +} +#if defined(STRATA_DS4_RUNTIME) +// A synchronising call is a barrier for BOTH queues: the dispatches deferred in mtl_rt's command buffer have to run before +// the host reads or overwrites what they touch (a window's last layers still pending while the next window's embedding rows +// are copied over the buffer its embed reads, measured: window 2's residual stream was window 1's stale tail). +// commit every stream's pending copies/fills, then wait for all of them (commit first: a stream may wait on another's event) +extern "C" void ds4_gpu_sync_all_streams() { + std::vector v; + { std::lock_guard lk(g_all_streams_mu); v = g_all_streams; } + for (StreamImpl* s : v) if (!s->rec) flush(s); + for (StreamImpl* s : v) if (!s->rec) wait_value(s->tl, s->submitted); +} +static thread_local bool t_in_barrier = false; +static void both_queues_barrier() { + if (t_in_barrier) return; + t_in_barrier = true; + ds4_metal_flush(); + t_in_barrier = false; +} + +#endif +Error memset_device(void* p, int value, size_t bytes) { +#if defined(STRATA_DS4_RUNTIME) + both_queues_barrier(); +#endif + StreamImpl* s = S(nullptr); + const Error e = encode_fill(s, p, value, bytes); + flush(s); + wait_value(s->tl, s->submitted); + return e; +} +Error memset_async(void* p, int value, size_t bytes, Stream st) { +#if defined(STRATA_DS4_RUNTIME) + if (!S(st)->rec && ctx_queue_ops() && bytes > 0) { + size_t o1 = 0; + if (find(p, &o1) && ds4_ctx_fill(p, value, bytes)) return kSuccess; + } + if (!S(st)->rec && !std::getenv("DS4_NO_QBARRIER")) both_queues_barrier(); +#endif + StreamImpl* s = S(st); + if (s->rec) { + Op o; + o.kind = Op::Fill; o.dst = p; o.value = value; o.bytes = bytes; + s->rec->ops.push_back(std::move(o)); + return kSuccess; + } + return encode_fill(s, p, value, bytes); +} + +Error memcpy_sync(void* dst, const void* src, size_t bytes, CopyKind) { +#if defined(STRATA_DS4_RUNTIME) + both_queues_barrier(); +#endif + StreamImpl* s = S(nullptr); + const Error e = encode_copy(s, dst, src, bytes); + flush(s); + wait_value(s->tl, s->submitted); + return e; +} +Error memcpy_2d_async(void* dst, size_t dpitch, const void* src, size_t spitch, size_t width, size_t height, CopyKind k, Stream st) { + for (size_t r = 0; r < height; ++r) { // skeleton: one copy per row (rare in the engine) + const Error e = memcpy_async((uint8_t*) dst + r * dpitch, (const uint8_t*) src + r * spitch, width, k, st); + if (e != kSuccess) return e; + } + return kSuccess; +} +Error memcpy_peer_async(void* dst, int dst_dev, const void* src, int src_dev, size_t bytes, Stream st) { + if (dst_dev != 0 || src_dev != 0) return fail(kErrorInvalidValue); // one GPU + return memcpy_async(dst, src, bytes, kDeviceToDevice, st); +} +Error memcpy_async(void* dst, const void* src, size_t bytes, CopyKind, Stream st) { +#if defined(STRATA_DS4_RUNTIME) + if (!S(st)->rec && ctx_queue_ops() && bytes > 0) { + size_t o1 = 0, o2 = 0; + if (find(dst, &o1) && find(src, &o2) && ds4_ctx_copy(dst, src, bytes)) return kSuccess; + } + if (!S(st)->rec && !std::getenv("DS4_NO_QBARRIER")) both_queues_barrier(); +#endif + StreamImpl* s = S(st); + if (s->rec) { + Op o; + o.kind = Op::Copy; o.dst = dst; o.src = src; o.bytes = bytes; + s->rec->ops.push_back(std::move(o)); + return kSuccess; + } + return encode_copy(s, dst, src, bytes); +} + +Error stream_create(Stream* s, unsigned) { *s = make_stream(); return kSuccess; } +Error stream_destroy(Stream s) { + if (!s) return kSuccess; + flush(s); + wait_value(s->tl, s->submitted); +#if defined(STRATA_DS4_RUNTIME) + { std::lock_guard lk(g_all_streams_mu); + g_all_streams.erase(std::remove(g_all_streams.begin(), g_all_streams.end(), s), g_all_streams.end()); } +#endif + s->tl->release(); + s->q->release(); + delete s; + return kSuccess; +} +Error stream_synchronize(Stream st) { +#if defined(STRATA_DS4_RUNTIME) + both_queues_barrier(); +#endif + StreamImpl* s = S(st); + flush(s); + return wait_value(s->tl, s->submitted); +} +Error stream_query(Stream st) { + StreamImpl* s = S(st); + flush(s); + return s->tl->signaledValue() >= s->submitted ? kSuccess : kNotReady; +} +Error stream_wait_event(Stream st, Event e, unsigned) { + StreamImpl* s = S(st); + if (s->rec) { + Op o; + o.kind = Op::Wait; o.ev = e; + s->rec->ops.push_back(std::move(o)); + return kSuccess; + } + return wait_event(s, e); +} +Error launch_host_func(Stream st, HostFn fn, void* arg) { + // skeleton: in stream order by completing the stream first (the engine's uses are rare; Phase 2 replaces this by a + // completion handler + a shared-event wait so the host thread does not block) + StreamImpl* s = S(st); + if (s->rec) return fail(kErrorUnsupported); + flush(s); + wait_value(s->tl, s->submitted); + fn(arg); + return kSuccess; +} + +Error event_create(Event* e, unsigned) { *e = new EventImpl; return kSuccess; } +Error event_destroy(Event e) { +#if defined(STRATA_DS4_RUNTIME) + if (e && e->token && ds4_ctx_drop) { ds4_ctx_drop(e->token); e->token = nullptr; } // (a mark nobody waited for holds its command buffer) +#endif + delete e; + return kSuccess; +} +Error event_record(Event e, Stream st) { +#if defined(STRATA_DS4_RUNTIME) + if (ctx_queue_ops() && !S(st)->rec) { // the work the event follows is in the context's command buffer: end it, wait later + if (e->token && ds4_ctx_drop) ds4_ctx_drop(e->token); // (recorded again before it was waited for: the earlier mark would leak its command buffer) + e->token = ds4_ctx_mark(); + e->ev = nullptr; + e->value = 0; + return kSuccess; + } +#endif + StreamImpl* s = S(st); + if (s->rec) { + Op o; + o.kind = Op::Record; o.ev = e; + s->rec->ops.push_back(std::move(o)); + return kSuccess; + } + return record_event(e, s); +} +#if defined(STRATA_DS4_RUNTIME) +Error event_query(Event e) { + if (e->token) return kNotReady; // (a mark waits for its wait) + return !e->ev || e->ev->signaledValue() >= e->value ? kSuccess : kNotReady; +} +Error event_synchronize(Event e) { + if (e->token) { + void* t = e->token; + e->token = nullptr; + ds4_ctx_wait(t); + return kSuccess; + } + return e->ev ? wait_value(e->ev, e->value) : kSuccess; +} +#else +Error event_query(Event e) { return !e->ev || e->ev->signaledValue() >= e->value ? kSuccess : kNotReady; } +Error event_synchronize(Event e) { return e->ev ? wait_value(e->ev, e->value) : kSuccess; } +#endif +Error event_elapsed_ms(float* ms, Event, Event) { *ms = 0.0f; return fail(kErrorUnsupported); } + +Error begin_capture(Stream st) { + StreamImpl* s = S(st); + if (s->rec) return fail(kErrorInvalidValue); + flush(s); + s->rec = new GraphImpl; + return kSuccess; +} +Error end_capture(Stream st, Graph* g) { + StreamImpl* s = S(st); + if (!s->rec) { *g = nullptr; return fail(kErrorInvalidValue); } + *g = s->rec; + s->rec = nullptr; + return kSuccess; +} +Error graph_node_count(Graph g, size_t* n) { *n = g ? g->ops.size() : 0; return kSuccess; } +Error graph_instantiate(GraphExec* x, Graph g) { + *x = new GraphExecImpl{g->ops}; + return kSuccess; +} +Error graph_launch(GraphExec x, Stream st) { + StreamImpl* s = S(st); + if (s->rec) { // a graph launched into a capture becomes part of it (CUDA: a child graph) + s->rec->ops.insert(s->rec->ops.end(), x->ops.begin(), x->ops.end()); + return kSuccess; + } + for (const Op& o : x->ops) { + const Error e = run_op(s, o); + if (e != kSuccess) return e; + } + return kSuccess; +} +Error graph_upload(GraphExec, Stream) { return kSuccess; } +Error graph_destroy(Graph g) { delete g; return kSuccess; } +Error graph_exec_destroy(GraphExec x) { delete x; return kSuccess; } + +// ---------------------------------------------------------------------------------------------------- metal.hpp +namespace metal { + +MTL::Device* device() { return rt().dev; } + +MTL::ComputePipelineState* pipeline(const char* source_name, const char* source, const char* function, std::string& err) { + Runtime& r = rt(); + const std::string key = std::string(source_name) + "::" + function; + std::lock_guard lk(r.mu); + if (auto it = r.pipes.find(key); it != r.pipes.end()) return it->second; + MTL::Library*& lib = r.libs[source_name]; + NS::Error* e = nullptr; + if (!lib) { + MTL::CompileOptions* co = MTL::CompileOptions::alloc()->init(); + co->setMathMode(MTL::MathModeSafe); // no fast-math reassociation: parity with the CUDA reference first + lib = r.dev->newLibrary(NS::String::string(source, NS::UTF8StringEncoding), co, &e); + co->release(); + if (!lib) { + err = std::string("MSL compile of ") + source_name + ": " + (e ? e->localizedDescription()->utf8String() : "?"); + return nullptr; + } + } + MTL::Function* f = lib->newFunction(NS::String::string(function, NS::UTF8StringEncoding)); + if (!f) { err = std::string("no MSL function ") + function; return nullptr; } + MTL::ComputePipelineState* p = r.dev->newComputePipelineState(f, &e); + f->release(); + if (!p) { err = std::string("pipeline ") + function + ": " + (e ? e->localizedDescription()->utf8String() : "?"); return nullptr; } + r.pipes[key] = p; + return p; +} + +Error dispatch(Stream st, MTL::ComputePipelineState* p, const void* args, size_t arg_bytes, Grid groups, Grid threads, + size_t tg_bytes) { + if (arg_bytes > 4096) return fail(kErrorInvalidValue); // setBytes limit + { // Metal skips a dispatch above the pipeline's maxTotalThreadsPerThreadgroup without any error (#2323): stop instead + const uint64_t want = (uint64_t) threads.x * threads.y * threads.z; + if (p && want > (uint64_t) p->maxTotalThreadsPerThreadgroup()) { + Runtime& r = rt(); + std::string name = "?"; + { + std::lock_guard lk(r.mu); + for (const auto& kv : r.pipes) + if (kv.second == p) { name = kv.first; break; } + } + std::fprintf(stderr, + "strata metal: the kernel %s can run at most %lu threads per threadgroup on this GPU (%s), but the " + "launch asks for %lu.\nMetal skips such a dispatch without any error, so every answer would be " + "wrong; stopping here instead. Please report this (the GPU, macOS version and this line) at " + "https://github.com/Niko1221/Strata/issues . (#2323)\n", + name.c_str(), (unsigned long) p->maxTotalThreadsPerThreadgroup(), + r.dev->name()->utf8String(), (unsigned long) want); + std::fflush(stderr); + std::exit(6); + } + } + StreamImpl* s = S(st); + Op o; + o.kind = Op::Dispatch; + o.pipe = p; + if (arg_bytes <= sizeof o.args) { std::memcpy(o.args, args, arg_bytes); o.arg_bytes = arg_bytes; } + else o.big_args.assign((const uint8_t*) args, (const uint8_t*) args + arg_bytes); + o.groups = groups; + o.threads = threads; + o.tg_bytes = tg_bytes; + if (s->rec) { s->rec->ops.push_back(std::move(o)); return kSuccess; } + return encode_dispatch(s, o); +} + +Stats stats() { return rt().stats; } + +} // namespace metal +} // namespace strata::gpu + +#if defined(STRATA_DS4_RUNTIME) +extern "C" void* ds4_gpu_buffer_for_cpu(const void* cpu, size_t bytes, size_t* off) { + return strata::gpu::buffer_for_cpu_impl(cpu, bytes, off); +} + +// The residency set every strata::gpu stream queue carries: the Context's queue (mtl_rt) needs it too, because kernels +// dereference raw gpuAddresses (weights, ring buffers, job arrays) that no setBuffer names, and memory that is not resident +// reads as zeros. +extern "C" void* ds4_gpu_residency_set() { return (void*) strata::gpu::residency_set_impl(); } + +// bytes from `cpu` to the end of the allocation that holds it (0: not an allocation) +extern "C" size_t ds4_gpu_alloc_remaining(const void* cpu) { + size_t off = 0; + void* b = strata::gpu::buffer_for_cpu_impl(cpu, 1, &off); + if (!b) return 0; + return ((MTL::Buffer*) b)->length() - off; +} + +#endif \ No newline at end of file diff --git a/edge0-strata/metal/src/mtl_rt.h b/edge0-strata/metal/src/mtl_rt.h new file mode 100644 index 00000000..fb1e6b63 --- /dev/null +++ b/edge0-strata/metal/src/mtl_rt.h @@ -0,0 +1,150 @@ +// metal/src/mtl_rt.h - a thin C++ face over Metal for the standalone kernel programs (parity, benches). +// +// The implementation (mtl_rt.mm) is the only Objective-C++ file; every program and test is plain C++ against this +// header. Kernels are MSL source files compiled at run time (newLibraryWithSource): the Command Line Tools have no +// offline `metal` compiler, and the runtime compiler is part of the OS. ggml-common.h (MIT, third_party/ggml) is +// prepended for the i-quant grids and block structs. +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace mtl { + +struct Buffer { + void* handle = nullptr; // id (retained); the first chunk of a chunked buffer + void* ptr = nullptr; // CPU address (shared storage) + size_t bytes = 0; + // a buffer longer than maxBufferLength: overlapping chunks every `stride` bytes (see Context::wrap) + std::shared_ptr> chunks; + size_t stride = 0; + size_t bias = 0; // bytes from the MTLBuffer's start to ptr (a view onto an allocation the runtime owns) + template T* as() const { return static_cast(ptr); } +}; + +struct Pipeline { + void* handle = nullptr; // id + int max_threads = 0; // maxTotalThreadsPerThreadgroup + int simd_width = 0; // threadExecutionWidth + int static_tg_mem = 0; // staticThreadgroupMemoryLength + std::string name; // the function (+ its constants), for messages +}; + +// A function-constant value set (index -> int / bool / uint). +struct Constants { + struct V { int index; int type; uint32_t v; }; // type 0 int, 1 bool, 2 uint, 3 short, 4 ushort + std::vector vals; + Constants& i(int idx, int v) { vals.push_back({idx, 0, (uint32_t)v}); return *this; } + Constants& b(int idx, bool v) { vals.push_back({idx, 1, (uint32_t)v}); return *this; } + Constants& u(int idx, uint32_t v) { vals.push_back({idx, 2, v}); return *this; } +}; + +struct Arg { // one argument of a dispatch: a buffer (+ offset) or inline bytes + const Buffer* buf = nullptr; + size_t off = 0; + std::vector bytes; + static Arg B(const Buffer& b, size_t off = 0) { Arg a; a.buf = &b; a.off = off; return a; } + template static Arg V(const T& v) { + Arg a; a.bytes.resize(sizeof(T)); std::memcpy(a.bytes.data(), &v, sizeof(T)); return a; + } +}; + +struct Dispatch { + const Pipeline* p = nullptr; + std::vector args; + uint32_t grid[3] = {1, 1, 1}; // threadgroups + uint32_t tg[3] = {32, 1, 1}; // threads per threadgroup + uint32_t tg_mem = 0; // dynamic threadgroup memory at index 0 (bytes) + const Buffer* indirect = nullptr; // grid from a GPU buffer (3 x uint32 threadgroup counts) at ind_off + size_t ind_off = 0; + bool barrier = true; // run_graph: a buffer memory barrier before this dispatch +}; + +// A pipeline whose maxTotalThreadsPerThreadgroup is below the threadgroup size a dispatch asks for is NOT executed by +// Metal and raises no error (0.1.44, #2323: gdn_step needs 512 threads, an M2 Max gave 448 for the GDN_NTOK >= 2 +// specializations - the layers silently kept their old outputs and every answer was garbage). So every dispatch is +// checked: too few threads is fatal, with a message that names the kernel, the limit and the size needed. +// STRATA_METAL_MAX_THREADS_CAP=N lowers every pipeline's reported limit to N (a test hook: it proves this check). +void require_fit(const Pipeline& p, uint32_t threads, const char* where); +inline void require_fit(const Dispatch& d, const char* where = "dispatch") { + require_fit(*d.p, d.tg[0] * d.tg[1] * d.tg[2], where); +} + +class Context { +public: + Context(); + ~Context(); + std::string device_name() const; + bool has_family_apple(int n) const; // MTLGPUFamilyApple + bool supports_metal4() const; + uint64_t recommended_working_set() const; + uint64_t max_buffer_length() const; + // Keeps the GPU out of its low-clock idle state while dispatches keep coming (a lone kernel in a quiet GPU runs 2-3x slower: + // measured 205 vs 69 us for the same matvec): a spin kernel of one threadgroup on a second queue, stopped 50 ms after the + // last dispatch. DS4_NO_KEEPALIVE=1: off. + void start_keepalive(); + // Zero-copy weights: the GPU address of host memory [p, p + bytes) (file-backed pages of the model, RAM slots) through a + // no-copy buffer over the pages it spans, made resident. 0: not mappable (the cap on buffers is reached, or the pages are not + // ours). The first dispatch that reads cold pages waits for them (the buffer is wired on use). + uint64_t zc_map(const void* p, size_t bytes); + Pipeline pipeline_uncached(int lib, const std::string& fn, const Constants& c); + size_t zc_count() const; + uint64_t allocated() const { return allocated_; } // bytes of alloc() + wrap() buffers so far + + // Compiles a library from MSL source files (paths), prepended by `prelude` (defines) and, if include_ggml, + // ggml-common.h with GGML_COMMON_DECL_METAL/IMPL_METAL. Returns a library id (index); aborts with the + // compiler log on error. + // Safe math unless STRATA_METAL_MATH=fast (the argument is ignored; kept for callers). + int compile(const std::vector& files, const std::string& prelude = "", bool fast_math = false); + Pipeline pipeline(int lib, const std::string& fn, const Constants& c = {}); + + Buffer alloc(size_t bytes); // shared storage, zeroed + Buffer wrap(void* page_aligned, size_t bytes); // no-copy wrap (page-aligned, size multiple of page) + // the same, also when bytes > maxBufferLength: then overlapping chunks, each holding any range of up to + // max_span bytes that starts in it (max_span = the largest tensor bound at one offset) + Buffer wrap(void* page_aligned, size_t bytes, size_t max_span); + void free(Buffer& b); + + // Encodes `ds` into ONE compute encoder (serial or concurrent dispatch type) and ONE command buffer, `reps` + // times over; commits and waits. Returns the GPU time (GPUEndTime - GPUStartTime) in seconds. + void* mark(); // split_runs() that returns a token for the buffer just committed (nullptr: nothing was queued) + void wait_mark(void* token); // block until that buffer has completed, then drop the token + void drop_mark(void* token); // drop a token without waiting (an event recorded again, or destroyed, before it was waited on) + void split_runs(); // commit the batched dispatches as their own command buffer, without waiting + void flush_runs(); // commit and wait for the dispatches batched since the last flush + double run(const std::vector& ds, int reps = 1, bool concurrent = false, bool barrier_between = false); + // A dependency graph in ONE concurrent encoder: a memory barrier only before dispatches with `barrier` set, so + // the dispatches between two barriers may overlap. Returns the GPU time in seconds. + double run_graph(const std::vector& ds); + // run_graph in segments: the dispatches [cuts[s-1], cuts[s]) are one concurrent encoder; after segment s (for + // s < cuts.size()) the GPU signals a shared event and waits until the host has run at_cut(s) - the host's turn + // between two parts of one committed command buffer (the SSD-streaming handshake: the routing is read, the + // missing experts are put in place). Returns the GPU time in seconds. + double run_graph_seg(const std::vector& ds, const std::vector& cuts, void (*at_cut)(void*, int), + void* user); + // One command buffer per `run_cb` call, empty or not: measures commit+wait round trips. + double run_wall(const std::vector& ds, int reps, double* gpu_seconds = nullptr); + + // Indirect command buffer: the dispatches recorded once, executed `reps` times in one command buffer. + // Buffers are bound per command (setKernelBuffer); inline bytes are not allowed (use buffers). + double run_icb(const std::vector& ds, int reps, bool barrier_between); + // Runs `ds` repeatedly for about `seconds` of wall time: lifts the GPU out of its idle clocks before a timing. + void warm(const std::vector& ds, double seconds = 0.3); + + std::string ggml_common_path; // where ggml-common.h is read from + std::string kernel_dir; // where relative kernel paths resolve + +private: + struct Impl; + Impl* impl_; + uint64_t allocated_ = 0; +}; + +double now_s(); // monotonic wall clock (seconds) +std::string read_file(const std::string& path); + +} // namespace mtl diff --git a/edge0-strata/metal/src/mtl_rt.mm b/edge0-strata/metal/src/mtl_rt.mm new file mode 100644 index 00000000..b78eab74 --- /dev/null +++ b/edge0-strata/metal/src/mtl_rt.mm @@ -0,0 +1,1117 @@ +// metal/src/mtl_rt.mm - see mtl_rt.h. Built with ARC (-fobjc-arc). +#import +#import + +#include "mtl_rt.h" +#include "resource_path.h" + +#include +#if defined(STRATA_DS4_RUNTIME) +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#endif +#include +#include +#include +#include + +namespace mtl { + +double now_s() { + return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +std::string read_file(const std::string& path) { + std::ifstream f(path, std::ios::binary); + if (!f) { + std::fprintf(stderr, "mtl: cannot read %s\n", path.c_str()); + std::exit(2); + } + std::stringstream ss; + ss << f.rdbuf(); + return ss.str(); +} + +#if defined(STRATA_DS4_RUNTIME) +extern "C" void* ds4_gpu_residency_set() __attribute__((weak_import)); +#endif +struct Context::Impl { +#if defined(STRATA_DS4_RUNTIME) + id cb_def = nil; // the open command buffer single dispatches accumulate in + id enc_def = nil; + int def_n = 0; + std::string def_names, inflight_names; // (DS4_ERRINFO) the kernels of the open / the committed command buffer + id qev = nil; // consecutive deferred buffers are chained through it (untracked buffers: the queue alone does not order them) + uint64_t qev_v = 0; + id inflight = nil; // the last committed-but-not-waited deferred buffer (the queue runs them in order) + std::vector graveyard; // buffers freed while the deferred command buffer still references them (it holds none) + std::recursive_mutex mu; // helper threads (the stager, the expert cache) flush too: one writer on the open encoder +#endif + id dev; + id q; + std::vector> libs; + id ev = nil; // run_graph_seg's handshake timeline + uint64_t ev_v = 0; +}; + +static inline id B(const Buffer& b) { return (__bridge id)b.handle; } +// the MTLBuffer and offset that bind byte `off` of b (a chunked buffer: the chunk the offset selects) +static inline id BO(const Buffer& b, size_t off, size_t& o) { + if (!b.chunks) { o = off + b.bias; return B(b); } // (bias: a view onto part of a bigger allocation; 0 otherwise) + const size_t k = std::min(off / b.stride, b.chunks->size() - 1); + o = off - k * b.stride; + return (__bridge id)(*b.chunks)[k]; +} +static inline id P(const Pipeline* p) { return (__bridge id)p->handle; } +static std::string g_device_name = "unknown GPU"; + +#if defined(STRATA_DS4_RUNTIME) + +// ---- DS4_TIMELINE=1: every deferred dispatch in its own encoder with start / end GPU timestamps; the report at exit gives, per kernel, +// its GPU time and the gap the GPU sat idle before it (host waits, event waits, launch). DS4_TL_FROM=N: only from dispatch N on. +static std::atomic g_disp_count{0}; // dispatches encoded +namespace { +struct Timeline { + bool on = false; + id tset = nil; + std::vector> bufs; // a sample buffer holds only ~1500 encoders: they rotate + size_t per = 0, cap = 0, n = 0; // encoders per buffer / in all (the buffers' limit) / recorded + long after = 0; // DS4_TL_AFTER: record from the Nth dispatch on + std::vector names; // per encoder + uint64_t cpu0 = 0, gpu0 = 0; + bool reg = false; + bool gate = true; // ds4_timeline_gate: record only while open + bool att = false; // the encoder just made carries a sample slot +} g_tl; + +id tl_make_encoder(id dev, id cb, MTLDispatchType dt) { + g_tl.att = false; + if (!g_tl.on) return [cb computeCommandEncoderWithDispatchType:dt]; + if (g_tl.tset == nil) { + for (id cs in dev.counterSets) + if ([cs.name isEqualToString:MTLCommonCounterSetTimestamp]) g_tl.tset = cs; + if (g_tl.tset == nil) { + std::fprintf(stderr, "mtl: DS4_TIMELINE: no timestamp counter set\n"); + g_tl.on = false; + return [cb computeCommandEncoderWithDispatchType:dt]; + } + for (size_t c = 4096; c >= 64 && g_tl.per == 0; c /= 2) { + MTLCounterSampleBufferDescriptor* d = [MTLCounterSampleBufferDescriptor new]; + d.counterSet = g_tl.tset; + d.storageMode = MTLStorageModeShared; + d.sampleCount = c * 2; + NSError* err = nil; + id b = [dev newCounterSampleBufferWithDescriptor:d error:&err]; + if (b != nil) { + g_tl.per = c; + g_tl.bufs.push_back(b); + } + } + if (g_tl.per == 0) { + std::fprintf(stderr, "mtl: DS4_TIMELINE: no counter sample buffer\n"); + g_tl.on = false; + return [cb computeCommandEncoderWithDispatchType:dt]; + } + g_tl.cap = g_tl.per * 120; + MTLTimestamp c0 = 0, g0 = 0; + [dev sampleTimestamps:&c0 gpuTimestamp:&g0]; + g_tl.cpu0 = c0; + g_tl.gpu0 = g0; + g_tl.names.assign(g_tl.cap, std::string()); + g_tl.after = std::getenv("DS4_TL_AFTER") ? std::atol(std::getenv("DS4_TL_AFTER")) : 0; + if (std::getenv("DS4_TL_DRAFT")) g_tl.gate = false; + } + if (!g_tl.gate || g_tl.n >= g_tl.cap || (long) g_disp_count.load() < g_tl.after) return [cb computeCommandEncoderWithDispatchType:dt]; + const size_t bi = g_tl.n / g_tl.per, si = g_tl.n % g_tl.per; + if (bi >= g_tl.bufs.size()) { + MTLCounterSampleBufferDescriptor* d = [MTLCounterSampleBufferDescriptor new]; + d.counterSet = g_tl.tset; + d.storageMode = MTLStorageModeShared; + d.sampleCount = g_tl.per * 2; + NSError* err = nil; + id b = [dev newCounterSampleBufferWithDescriptor:d error:&err]; + if (b == nil) { + g_tl.cap = g_tl.n; + return [cb computeCommandEncoderWithDispatchType:dt]; + } + g_tl.bufs.push_back(b); + } + MTLComputePassDescriptor* pd = [MTLComputePassDescriptor computePassDescriptor]; + pd.dispatchType = dt; + pd.sampleBufferAttachments[0].sampleBuffer = g_tl.bufs[bi]; + pd.sampleBufferAttachments[0].startOfEncoderSampleIndex = 2 * si; + pd.sampleBufferAttachments[0].endOfEncoderSampleIndex = 2 * si + 1; + g_tl.att = true; + return [cb computeCommandEncoderWithDescriptor:pd]; +} + +void tl_report(id dev) { + if (!g_tl.on && g_tl.bufs.empty()) return; + const size_t n = std::min(g_tl.n, g_tl.cap); + if (n < 2) return; + std::vector tsv(2 * n); + for (size_t b = 0; b * g_tl.per < n; ++b) { + const size_t cnt = std::min(g_tl.per, n - b * g_tl.per); + NSData* data = [g_tl.bufs[b] resolveCounterRange:NSMakeRange(0, 2 * cnt)]; + if (data == nil) { std::fprintf(stderr, "mtl: DS4_TIMELINE: resolve failed\n"); return; } + std::memcpy(&tsv[2 * b * g_tl.per], [data bytes], 2 * cnt * sizeof(MTLCounterResultTimestamp)); + } + const MTLCounterResultTimestamp* ts = tsv.data(); + MTLTimestamp c1 = 0, g1 = 0; + [dev sampleTimestamps:&c1 gpuTimestamp:&g1]; + mach_timebase_info_data_t tb; + mach_timebase_info(&tb); + const double cpu_ns_per_tick = (double) tb.numer / (double) tb.denom; + (void) cpu_ns_per_tick; // (sampleTimestamps' CPU value is in ns already) + const double ns_per_gpu = (g1 > g_tl.gpu0) ? (double) (c1 - g_tl.cpu0) / (double) (g1 - g_tl.gpu0) : 1.0; + std::fprintf(stderr, "timeline: cap %zu n %zu scale %.6f ns/tick first ts %llu %llu last %llu\n", g_tl.cap, g_tl.n, ns_per_gpu, (unsigned long long) ts[0].timestamp, (unsigned long long) ts[1].timestamp, (unsigned long long) ts[2 * n - 1].timestamp); + size_t from = std::getenv("DS4_TL_FROM") ? (size_t) std::atoll(std::getenv("DS4_TL_FROM")) : 0; + struct Acc { double busy = 0, gap = 0, gap_big = 0; size_t n = 0, n_big = 0; }; + std::map acc; + double t_first = 0, t_last = 0, total_busy = 0, total_gap = 0, total_big = 0; + double prev_end = -1, ovl_us = 0; + size_t used = 0, n_ovl = 0; + for (size_t k = from; k < n; ++k) { + if (g_tl.names[k].empty()) continue; + const double st = (double) ts[2 * k].timestamp * ns_per_gpu * 1e-3, en = (double) ts[2 * k + 1].timestamp * ns_per_gpu * 1e-3; // us + if (!(en >= st)) continue; + if (prev_end < 0) t_first = st; + Acc& a = acc[g_tl.names[k]]; + a.busy += en - st; + ++a.n; + total_busy += en - st; + if (prev_end >= 0 && st < prev_end) { ++n_ovl; ovl_us += prev_end - st; } + if (prev_end >= 0 && st >= prev_end) { + const double gp = st - prev_end; + a.gap += gp; + total_gap += gp; + if (gp > 300.0) { a.gap_big += gp; ++a.n_big; total_big += gp; } + } + prev_end = en; + t_last = en; + ++used; + } + std::vector> v; + for (auto& kv : acc) v.emplace_back(kv.second.busy + kv.second.gap, kv.first); + std::sort(v.begin(), v.end(), [](auto& x, auto& y) { return x.first > y.first; }); + std::fprintf(stderr, "timeline: %zu dispatches, wall %.1f ms, GPU busy %.1f ms, idle gaps %.1f ms (of it > 300 us: %.1f ms)\n", used, + (t_last - t_first) / 1000.0, total_busy / 1000.0, total_gap / 1000.0, total_big / 1000.0); + std::fprintf(stderr, "timeline: %zu dispatches started before the previous one ended (%.1f ms overlapped)\n", n_ovl, ovl_us / 1000.0); + std::fprintf(stderr, "timeline: kernel n busy ms us/disp gap ms gap/disp big-gap ms (n)\n"); + for (size_t i = 0; i < v.size() && i < 40; ++i) { + const Acc& a = acc[v[i].second]; + std::fprintf(stderr, "timeline: %-28s %6zu %9.1f %8.1f %8.1f %8.1f %10.1f (%zu)\n", v[i].second.c_str(), a.n, a.busy / 1000.0, a.busy / (double) a.n, + a.gap / 1000.0, a.gap / (double) a.n, a.gap_big / 1000.0, a.n_big); + } +} +} // namespace + +extern "C" void ds4_timeline_gate(int open) { g_tl.gate = open != 0; } + + +// ---- ds4_busy_begin / ds4_busy_end: a one-threadgroup kernel on a second queue that spins on a shared flag while the host works +// (the CPU lane), so the GPU is not idle between two bursts (idle -> the next kernels run 2-3x slower); the host sets the flag the +// moment it is done and the kernel leaves within microseconds. DS4_BUSY=0 turns it off. +namespace { +struct Busy { + bool init = false, ok = false; + id q = nil; + id ps = nil; + id flag = nil; + id last = nil; + bool running = false; +}; +Busy g_busy; +} // namespace + +extern "C" void ds4_busy_begin() { + static const bool off = std::getenv("DS4_BUSY") && std::string(std::getenv("DS4_BUSY")) == "0"; + if (off) return; + @autoreleasepool { + if (!g_busy.init) { + g_busy.init = true; + id dev = MTLCreateSystemDefaultDevice(); + static const char* kSrc = R"MSL( +#include +using namespace metal; +kernel void busy_spin(device volatile atomic_uint* flag [[buffer(0)]], constant uint& cap [[buffer(1)]], uint t [[thread_position_in_grid]]) { + uint x = t; + float f = (float) t * 1e-9f; + for (uint i = 0; i < cap; ++i) { + for (uint j = 0; j < 64; ++j) x = x * 1664525u + 1013904223u; + for (uint j = 0; j < 16; ++j) f = fma(f, 1.0000001f, 1e-7f); // (a real FP chain: keeps the ALU clock up) + x ^= as_type(f); + if (atomic_load_explicit(flag, memory_order_relaxed) != 0u) break; + } + if (x == 0xdeadbeefu) atomic_store_explicit(flag, 2u, memory_order_relaxed); +} +)MSL"; + NSError* err = nil; + id lib = [dev newLibraryWithSource:[NSString stringWithUTF8String:kSrc] options:nil error:&err]; + if (lib) g_busy.ps = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"busy_spin"] error:&err]; + if (g_busy.ps) { + g_busy.q = [dev newCommandQueue]; + g_busy.flag = [dev newBufferWithLength:64 options:MTLResourceStorageModeShared]; + g_busy.ok = true; + } + } + if (!g_busy.ok || g_busy.running) return; + g_busy.running = true; + if (g_busy.last != nil) [g_busy.last waitUntilCompleted]; // (the previous spin has left) + *(volatile uint32_t*) [g_busy.flag contents] = 0; + uint32_t cap = std::getenv("DS4_BUSY_CAP") ? (uint32_t) std::atoi(std::getenv("DS4_BUSY_CAP")) : 20000u; // ~ tens of ms at most + // the spin's width: one 256-thread threadgroup leaves most of the GPU idle and the clock drops anyway; DS4_BUSY_W + // threadgroups (256 threads each) keep more of the card's ALUs fed, which is what holds the performance state. + static const uint32_t width = std::getenv("DS4_BUSY_W") ? (uint32_t) std::max(1, std::atoi(std::getenv("DS4_BUSY_W"))) : 1u; + id cb = [g_busy.q commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:g_busy.ps]; + [enc setBuffer:g_busy.flag offset:0 atIndex:0]; + [enc setBytes:&cap length:4 atIndex:1]; + [enc dispatchThreads:MTLSizeMake(256 * width, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + [enc endEncoding]; + [cb commit]; + g_busy.last = cb; + } +} +extern "C" void ds4_busy_end() { + if (!g_busy.ok || g_busy.last == nil || !g_busy.running) return; + *(volatile uint32_t*) [g_busy.flag contents] = 1; + g_busy.running = false; +} + +#endif +Context::Context() : impl_(new Impl) { + impl_->dev = MTLCreateSystemDefaultDevice(); +#if defined(STRATA_DS4_RUNTIME) + if (std::getenv("DS4_TIMELINE") && !g_tl.reg) { + g_tl.on = true; + g_tl.reg = true; + static id tl_dev = impl_->dev; + std::atexit([] { tl_report(tl_dev); }); + } +#endif + if (!impl_->dev) { + std::fprintf(stderr, "mtl: no Metal device\n"); + std::exit(2); + } + impl_->q = [impl_->dev newCommandQueue]; +#if defined(STRATA_DS4_RUNTIME) + if (ds4_gpu_residency_set && !std::getenv("DS4_NO_RESIDENCY")) { + id rs = (__bridge id)ds4_gpu_residency_set(); + if (rs) [impl_->q addResidencySet:rs]; + } +#endif + // env, then /metal/ (a component package), then the compiled-in source-tree path + ggml_common_path = strata_metal::resource("STRATA_GGML_COMMON", "ggml-common.h", STRATA_GGML_COMMON_DEFAULT); + kernel_dir = strata_metal::resource("STRATA_METAL_KERNELS", "kernels", STRATA_METAL_KERNELS_DEFAULT); +} + +Context::~Context() { delete impl_; } + +std::string Context::device_name() const { return [[impl_->dev name] UTF8String]; } +bool Context::has_family_apple(int n) const { + return [impl_->dev supportsFamily:(MTLGPUFamily)(MTLGPUFamilyApple1 + (n - 1))]; +} +bool Context::supports_metal4() const { + if (@available(macOS 26.0, *)) return [impl_->dev supportsFamily:MTLGPUFamilyMetal4]; + return false; +} +uint64_t Context::recommended_working_set() const { return [impl_->dev recommendedMaxWorkingSetSize]; } +uint64_t Context::max_buffer_length() const { return [impl_->dev maxBufferLength]; } + +int Context::compile(const std::vector& files, const std::string& prelude, bool fast_math) { + std::string src = prelude + "\n#include \nusing namespace metal;\n"; + if (!ggml_common_path.empty() && prelude.find("NO_GGML_COMMON") == std::string::npos) { + src += "#define GGML_COMMON_DECL_METAL\n#define GGML_COMMON_IMPL_METAL\n"; + src += read_file(ggml_common_path); + src += "\n"; + } + for (const auto& f : files) { + std::string p = (f.size() && f[0] == '/') ? f : kernel_dir + "/" + f; + src += "#line 1 \"" + f + "\"\n" + read_file(p) + "\n"; + } + MTLCompileOptions* o = [MTLCompileOptions new]; + // Safe math by default: no reassociation, so every token of a window is bitwise what a one-token call gives + // (the kernels write their contractions as explicit fma). Measured free on the M5 Max. + fast_math = false; + if (const char* m = std::getenv("STRATA_METAL_MATH")) fast_math = std::string(m) == "fast"; + if (@available(macOS 15.0, *)) { + o.mathMode = fast_math ? MTLMathModeFast : MTLMathModeSafe; + } else { + o.fastMathEnabled = fast_math; + } + NSError* err = nil; + id lib = [impl_->dev newLibraryWithSource:[NSString stringWithUTF8String:src.c_str()] options:o error:&err]; + if (!lib) { + std::fprintf(stderr, "mtl: compile failed (%s, macOS %s; %s):\n%s\n", [[impl_->dev name] UTF8String], + [[[NSProcessInfo processInfo] operatingSystemVersionString] UTF8String], + files.empty() ? "?" : files.back().c_str(), [[err localizedDescription] UTF8String]); + std::exit(3); + } + impl_->libs.push_back(lib); + return (int)impl_->libs.size() - 1; +} + +#if defined(STRATA_DS4_RUNTIME) + +// ---- the keep-alive spin (see mtl_rt.h) +static std::atomic g_last_act_ns{0}; +static inline int64_t ka_now_ns() { return std::chrono::duration_cast(std::chrono::steady_clock::now().time_since_epoch()).count(); } +static std::atomic g_ka_on{false}; + +void Context::start_keepalive() { + if (!std::getenv("DS4_KEEPALIVE") || g_ka_on.exchange(true)) return; // (measured slower in the decode loop) + id dev = impl_->dev; + std::thread([dev]() { + @autoreleasepool { + static const char* kSpin = R"MSL( +#include +using namespace metal; +kernel void ka_spin(device atomic_uint* c [[buffer(0)]], constant uint& n [[buffer(1)]], uint t [[thread_position_in_grid]]) { + uint x = t; + for (uint i = 0; i < n; ++i) x = x * 1664525u + 1013904223u; + if (x == 0xdeadbeefu) atomic_fetch_add_explicit(c, 1u, memory_order_relaxed); +} +)MSL"; + NSError* err = nil; + id lib = [dev newLibraryWithSource:[NSString stringWithUTF8String:kSpin] options:nil error:&err]; + if (!lib) return; + id ps = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"ka_spin"] error:&err]; + if (!ps) return; + id q = [dev newCommandQueue]; + id cbuf = [dev newBufferWithLength:64 options:MTLResourceStorageModeShared]; + uint32_t iters = std::getenv("DS4_KEEPALIVE_ITERS") ? (uint32_t) std::atoi(std::getenv("DS4_KEEPALIVE_ITERS")) : 20000u; + for (;;) { + if (ka_now_ns() - g_last_act_ns.load(std::memory_order_relaxed) > 50000000LL) { + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + continue; + } + @autoreleasepool { + id cb = [q commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:ps]; + [enc setBuffer:cbuf offset:0 atIndex:0]; + [enc setBytes:&iters length:4 atIndex:1]; + [enc dispatchThreads:MTLSizeMake(256, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + } + } + } + }).detach(); +} + + +// ---- zero-copy weights (see mtl_rt.h) +namespace { +struct ZcBuf { size_t len; void* handle; uint64_t addr; }; +std::map g_zc; +std::mutex g_zc_mu; +} // namespace + +uint64_t Context::zc_map(const void* p, size_t bytes) { + const size_t pg = 16384; + const uintptr_t lo = (uintptr_t) p & ~(uintptr_t) (pg - 1); + const uintptr_t hi = ((uintptr_t) p + bytes + pg - 1) & ~(uintptr_t) (pg - 1); + std::lock_guard lk(g_zc_mu); + auto it = g_zc.upper_bound(lo); + if (it != g_zc.begin()) { + --it; + if (it->first <= lo && hi <= it->first + it->second.len) return it->second.addr + ((uintptr_t) p - it->first); + } + static const size_t cap = std::getenv("DS4_ZC_MAX") ? (size_t) std::atoll(std::getenv("DS4_ZC_MAX")) : 9000; + if (g_zc.size() >= cap) return 0; + id b = [impl_->dev newBufferWithBytesNoCopy:(void*) lo length:hi - lo options:MTLResourceStorageModeShared deallocator:nil]; + if (!b) return 0; + if (ds4_gpu_residency_set) { + id rs = (__bridge id)ds4_gpu_residency_set(); + if (rs) { + [rs addAllocation:b]; + [rs commit]; + } + } + ZcBuf z{hi - lo, (__bridge_retained void*) b, (uint64_t) [b gpuAddress]}; + g_zc[lo] = z; + return z.addr + ((uintptr_t) p - lo); +} +size_t Context::zc_count() const { std::lock_guard lk(g_zc_mu); return g_zc.size(); } + +#endif +Pipeline Context::pipeline(int lib, const std::string& fn, const Constants& c) { +#if defined(STRATA_DS4_RUNTIME) + // (cached: the host helpers ask for the same pipeline a dispatch at a time; newComputePipelineState costs ~0.3 ms and + // [MTLDevice name] reads the preferences - the prompt path's expert copies (3 a streamed expert) spent seconds there) + static std::mutex cache_mu; + static std::map cache; + std::string key = std::to_string(lib) + "|" + fn; + for (const auto& v : c.vals) key += "|" + std::to_string(v.index) + ":" + std::to_string(v.type) + ":" + std::to_string(v.v); + key += "|" + std::to_string((uintptr_t) this); + { + std::lock_guard lk(cache_mu); + auto it = cache.find(key); + if (it != cache.end()) return it->second; + } + Pipeline made = pipeline_uncached(lib, fn, c); + { + std::lock_guard lk(cache_mu); + cache[key] = made; + } + return made; +} + +Pipeline Context::pipeline_uncached(int lib, const std::string& fn, const Constants& c) { +#endif + NSError* err = nil; + id f; + NSString* name = [NSString stringWithUTF8String:fn.c_str()]; + if (c.vals.empty()) { + f = [impl_->libs[lib] newFunctionWithName:name]; + // a function with (optional) function constants needs a specialization even when none is set + if (f && f.functionConstantsDictionary.count > 0) + f = [impl_->libs[lib] newFunctionWithName:name constantValues:[MTLFunctionConstantValues new] error:&err]; + } else { + MTLFunctionConstantValues* cv = [MTLFunctionConstantValues new]; + for (const auto& v : c.vals) { + int32_t iv = (int32_t)v.v; + bool bv = v.v != 0; + uint32_t uv = v.v; + switch (v.type) { + case 0: [cv setConstantValue:&iv type:MTLDataTypeInt atIndex:v.index]; break; + case 1: [cv setConstantValue:&bv type:MTLDataTypeBool atIndex:v.index]; break; + default: [cv setConstantValue:&uv type:MTLDataTypeUInt atIndex:v.index]; break; + } + } + f = [impl_->libs[lib] newFunctionWithName:name constantValues:cv error:&err]; + } + if (!f) { + std::fprintf(stderr, "mtl: no function %s %s\n", fn.c_str(), err ? [[err localizedDescription] UTF8String] : ""); + std::exit(3); + } + MTLComputePipelineDescriptor* d = [MTLComputePipelineDescriptor new]; + d.computeFunction = f; + d.supportIndirectCommandBuffers = YES; + id ps = [impl_->dev newComputePipelineStateWithDescriptor:d + options:MTLPipelineOptionNone + reflection:nil + error:&err]; + if (!ps) { + std::fprintf(stderr, "mtl: pipeline %s: %s\n", fn.c_str(), [[err localizedDescription] UTF8String]); + std::exit(3); + } + Pipeline p; + p.handle = (__bridge_retained void*)ps; + p.max_threads = (int)ps.maxTotalThreadsPerThreadgroup; + if (const char* cap = std::getenv("STRATA_METAL_MAX_THREADS_CAP")) { // test hook, see require_fit + const int lim = std::atoi(cap); + if (lim > 0 && lim < p.max_threads) p.max_threads = lim; + } + p.simd_width = (int)ps.threadExecutionWidth; + p.static_tg_mem = (int)ps.staticThreadgroupMemoryLength; + p.name = fn; + if (!c.vals.empty()) { + p.name += "{"; + for (size_t i = 0; i < c.vals.size(); ++i) + p.name += (i ? "," : "") + std::to_string(c.vals[i].index) + "=" + std::to_string((int)c.vals[i].v); + p.name += "}"; + } +#if defined(STRATA_DS4_RUNTIME) + if (g_device_name == "unknown GPU") g_device_name = [[impl_->dev name] UTF8String]; +#else + g_device_name = [[impl_->dev name] UTF8String]; +#endif + return p; +} + +void require_fit(const Pipeline& p, uint32_t threads, const char* where) { + if (threads == 0 || (int64_t)threads <= (int64_t)p.max_threads) return; + std::fprintf(stderr, + "strata metal: the kernel %s can run at most %d threads per threadgroup on this GPU (%s), but the %s asks " + "for %u.\nMetal skips such a dispatch without any error, so every answer would be wrong; stopping here " + "instead. Please report this (the GPU, macOS version and this line) at " + "https://github.com/Niko1221/Strata/issues . (#2323)\n", + p.name.empty() ? "?" : p.name.c_str(), p.max_threads, g_device_name.c_str(), where, threads); + std::fflush(stderr); + std::exit(6); +} + +Buffer Context::alloc(size_t bytes) { +#if defined(STRATA_DS4_RUNTIME) + // Metal returns the new buffer with one more reference parked in the calling thread's autorelease pool. The engine's host + // threads have no pool (only Context::run opens one), so that reference was never given back: every buffer made here lived for + // the rest of the process, free() or not (a 16 KB job array per expert matvec, ~90K a request: +1.3 GB and a slower queue every + // call, 16 -> 3 tok/s after a few requests on a big Mac). + Buffer r; + @autoreleasepool { + id b = [impl_->dev newBufferWithLength:bytes options:MTLResourceStorageModeShared]; + if (!b) { + std::fprintf(stderr, "mtl: alloc %zu failed\n", bytes); + std::exit(4); + } + std::memset([b contents], 0, bytes); + r.handle = (__bridge_retained void*)b; + r.ptr = [b contents]; + } + allocated_ += bytes; + r.bytes = bytes; + return r; +#else + id b = [impl_->dev newBufferWithLength:bytes options:MTLResourceStorageModeShared]; + if (!b) { + std::fprintf(stderr, "mtl: alloc %zu failed\n", bytes); + std::exit(4); + } + std::memset([b contents], 0, bytes); + allocated_ += bytes; + Buffer r; + r.handle = (__bridge_retained void*)b; + r.ptr = [b contents]; + r.bytes = bytes; + return r; +#endif +} + +Buffer Context::wrap(void* p, size_t bytes, size_t max_span) { + const size_t pg = 16384, maxlen = [impl_->dev maxBufferLength] / pg * pg; + if (bytes <= maxlen || std::getenv("STRATA_METAL_WRAP_SPLIT") && std::string(std::getenv("STRATA_METAL_WRAP_SPLIT")) == "0") + return wrap(p, bytes); + const size_t span = (max_span + pg - 1) / pg * pg; + if (span >= maxlen / 2) { + std::fprintf(stderr, "mtl: wrap %zu: a tensor of %zu bytes does not fit the chunks\n", bytes, max_span); + std::exit(4); + } + Buffer r; + r.chunks = std::make_shared>(); + r.stride = maxlen - span; + for (size_t base = 0; base < bytes; base += r.stride) { + const size_t len = std::min(maxlen, (bytes - base + pg - 1) / pg * pg); + id b = [impl_->dev newBufferWithBytesNoCopy:(uint8_t*) p + base length:len + options:MTLResourceStorageModeShared deallocator:nil]; + if (!b) { + std::fprintf(stderr, "mtl: wrap chunk %zu+%zu failed\n", base, len); + std::exit(4); + } + r.chunks->push_back((__bridge_retained void*)b); + if (base + len >= bytes) break; + } + allocated_ += bytes; + r.handle = (*r.chunks)[0]; + r.ptr = p; + r.bytes = bytes; + std::fprintf(stderr, "mtl: wrapped %.2f GB as %zu overlapping buffers (maxBufferLength %.2f GB, overlap %.2f GB)\n", + bytes / 1e9, r.chunks->size(), maxlen / 1e9, span / 1e9); + return r; +} + +#if defined(STRATA_DS4_RUNTIME) +// The runtime's own MTLBuffer under a CPU address (strata::gpu's allocations), when [p, p + bytes) lies inside one. Binding it +// - the same MTLBuffer, so the same GPU address, as every other kernel that touches this memory - instead of a fresh no-copy +// buffer per call. (Each fresh wrapper is a second GPU mapping of the same pages: the quantiser read the activations through +// one mapping while the norm that produced them had written through another, and saw the previous layer's.) +extern "C" void* ds4_gpu_buffer_for_cpu(const void* cpu, size_t bytes, size_t* off) __attribute__((weak_import)); +extern "C" size_t ds4_gpu_alloc_remaining(const void* cpu) __attribute__((weak_import)); +#endif +Buffer Context::wrap(void* p, size_t bytes) { +#if defined(STRATA_DS4_RUNTIME) + if (ds4_gpu_buffer_for_cpu && !std::getenv("DS4_WRAP_COPY")) { + size_t off = 0; + if (void* mb = ds4_gpu_buffer_for_cpu(p, bytes, &off)) { + Buffer r; + r.handle = (__bridge_retained void*)((__bridge id)mb); + r.ptr = p; + r.bytes = bytes; + r.bias = off; + return r; + } + } +#endif +#if defined(STRATA_DS4_RUNTIME) + if (std::getenv("DS4_WRAP_LOG")) // debugging: a wrap that is not inside one runtime allocation gets a second GPU mapping of the pages + std::fprintf(stderr, "mtl: wrap FALLBACK p=%p bytes=%zu alloc_remaining=%zu\n", p, bytes, ds4_gpu_alloc_remaining ? ds4_gpu_alloc_remaining(p) : (size_t) 0); +#endif +#if defined(STRATA_DS4_RUNTIME) + Buffer r; + @autoreleasepool { // (see alloc: the engine's host threads have no pool of their own) + id b = [impl_->dev newBufferWithBytesNoCopy:p length:bytes options:MTLResourceStorageModeShared deallocator:nil]; + if (!b) { + std::fprintf(stderr, "mtl: wrap %zu failed\n", bytes); + std::exit(4); + } + r.handle = (__bridge_retained void*)b; + } + allocated_ += bytes; + r.ptr = p; + r.bytes = bytes; + return r; +#else + id b = [impl_->dev newBufferWithBytesNoCopy:p length:bytes options:MTLResourceStorageModeShared deallocator:nil]; + if (!b) { + std::fprintf(stderr, "mtl: wrap %zu failed\n", bytes); + std::exit(4); + } + allocated_ += bytes; + Buffer r; + r.handle = (__bridge_retained void*)b; + r.ptr = p; + r.bytes = bytes; + return r; +#endif +} + +void Context::free(Buffer& b) { +#if defined(STRATA_DS4_RUNTIME) + // The deferred command buffer is made WITHOUT retained references: a buffer released here while a dispatch that names it + // is still queued is read after it is gone (or after another alloc took its memory: the job arrays of iq_mv read the + // NEXT call's jobs, measured as rows of zeros and different tokens every run). Hold it until the buffer has completed. + std::lock_guard lk_(impl_->mu); + const bool hold = impl_->cb_def != nil || impl_->inflight != nil; +#endif + if (b.chunks) { +#if defined(STRATA_DS4_RUNTIME) + for (void* h : *b.chunks) { + if (hold) impl_->graveyard.push_back(h); + else { id x = (__bridge_transfer id)h; x = nil; } + } +#else + for (void* h : *b.chunks) { id x = (__bridge_transfer id)h; x = nil; } +#endif + b = Buffer{}; + return; + } + if (b.handle) { +#if defined(STRATA_DS4_RUNTIME) + if (hold) impl_->graveyard.push_back(b.handle); + else { id x = (__bridge_transfer id)b.handle; x = nil; } +#else + id x = (__bridge_transfer id)b.handle; + x = nil; +#endif + } + b = Buffer{}; +} + +#if defined(STRATA_DS4_RUNTIME) +static std::mutex g_hist_mu; +static std::unordered_map g_pipe_names; +static std::unordered_map g_kernel_counts; +// DS4_ENC_TIMES: where a dispatch's encoding time goes - the pipeline state, the arguments, the dispatch call +static std::atomic g_t_pipe{0}, g_t_args{0}, g_t_disp{0}; +#endif +static void encode(id enc, const Dispatch& d) { + require_fit(d, "dispatch"); +#if defined(STRATA_DS4_RUNTIME) + static const bool enc_times = std::getenv("DS4_ENC_TIMES") != nullptr; + const auto t_a = std::chrono::steady_clock::now(); +#endif + [enc setComputePipelineState:P(d.p)]; +#if defined(STRATA_DS4_RUNTIME) + const auto t_b = std::chrono::steady_clock::now(); +#endif + for (size_t i = 0; i < d.args.size(); ++i) { + const Arg& a = d.args[i]; + if (a.buf) { + size_t o; + id mb = BO(*a.buf, a.off, o); +#if defined(STRATA_DS4_RUNTIME) + static const bool chkbuf = std::getenv("DS4_CHECK_BUF") != nullptr; + if (chkbuf && o >= [mb length]) std::fprintf(stderr, "mtl: %s arg %zu: offset %zu past the buffer (%zu bytes; Buffer.bytes %zu bias %zu off %zu)\n", d.p->name.c_str(), i, o, (size_t) [mb length], a.buf->bytes, a.buf->bias, (size_t) a.off); +#endif + [enc setBuffer:mb offset:o atIndex:i]; + } + else + [enc setBytes:a.bytes.data() length:a.bytes.size() atIndex:i]; + } +#if defined(STRATA_DS4_RUNTIME) + const auto t_c = std::chrono::steady_clock::now(); +#endif + if (d.tg_mem) [enc setThreadgroupMemoryLength:(d.tg_mem + 15) / 16 * 16 atIndex:0]; + if (d.indirect) + [enc dispatchThreadgroupsWithIndirectBuffer:B(*d.indirect) + indirectBufferOffset:d.ind_off + threadsPerThreadgroup:MTLSizeMake(d.tg[0], d.tg[1], d.tg[2])]; + else + [enc dispatchThreadgroups:MTLSizeMake(d.grid[0], d.grid[1], d.grid[2]) + threadsPerThreadgroup:MTLSizeMake(d.tg[0], d.tg[1], d.tg[2])]; +#if defined(STRATA_DS4_RUNTIME) + if (enc_times) { + const auto t_d = std::chrono::steady_clock::now(); + g_t_pipe.fetch_add((uint64_t) std::chrono::duration_cast(t_b - t_a).count(), std::memory_order_relaxed); + g_t_args.fetch_add((uint64_t) std::chrono::duration_cast(t_c - t_b).count(), std::memory_order_relaxed); + g_t_disp.fetch_add((uint64_t) std::chrono::duration_cast(t_d - t_c).count(), std::memory_order_relaxed); + } +} + +static std::deque g_recent; // the last dispatches encoded (printed with a command buffer failure) +static std::string recent_str() { std::string r; for (auto& n : g_recent) r += n + " "; return r; } +static double g_gpu_ms = 0.0; +static std::atomic g_gpu_busy_us{0}; // DS4_GPU_TOTAL: the GPU time of every committed command buffer +static void add_busy_handler(id cb) { + static const bool on = std::getenv("DS4_GPU_TOTAL") != nullptr; + if (on) [cb addCompletedHandler:^(id c) { const double d = c.GPUEndTime - c.GPUStartTime; if (d > 0) g_gpu_busy_us += (uint64_t) (d * 1e6); }]; +} // GPU time of every committed command buffer, summed +static std::atomic g_wait_ns{0}, g_enc_ns{0}, g_wait_n{0}; // host time: waiting on command buffers; encoding dispatches +static std::mutex g_prof_mu; +static std::unordered_map> g_prof; // DS4_PROF: kernel -> (GPU us, dispatches) +static std::string g_prof_last; +static std::atomic g_cb_count{0}; +static bool g_count_reg = false; +static void ds4_report_counts() { + if (g_gpu_busy_us.load()) std::fprintf(stderr, "mtl: GPU busy (every command buffer) %.1f ms\n", g_gpu_busy_us.load() / 1000.0); + std::fprintf(stderr, "mtl: %.1f ms GPU, %llu dispatches, %llu command buffers\n", g_gpu_ms, + (unsigned long long) g_disp_count.load(), (unsigned long long) g_cb_count.load()); + std::fprintf(stderr, "mtl: host waited %.1f ms for the GPU in %llu waits; encoding the dispatches took %.1f ms\n", g_wait_ns.load() / 1e6, (unsigned long long) g_wait_n.load(), g_enc_ns.load() / 1e6); + if (g_t_pipe.load() || g_t_args.load() || g_t_disp.load()) + std::fprintf(stderr, "mtl: of the encoding: pipeline %.1f ms, arguments %.1f ms, dispatch %.1f ms\n", + g_t_pipe.load() / 1e6, g_t_args.load() / 1e6, g_t_disp.load() / 1e6); + if (!g_prof.empty()) { + std::vector> pv; + double tot = 0; + for (auto& kv : g_prof) { pv.emplace_back(kv.second.first, kv.first); tot += kv.second.first; } + std::sort(pv.begin(), pv.end(), [](auto& a, auto& b) { return a.first > b.first; }); + std::fprintf(stderr, "prof: kernel GPU ms n us/dispatch (total %.1f ms)\n", tot / 1000.0); + for (size_t i = 0; i < pv.size() && i < 30; ++i) { + auto& e = g_prof[pv[i].second]; + std::fprintf(stderr, "prof: %-30s %8.1f %6llu %8.1f\n", pv[i].second.c_str(), e.first / 1000.0, (unsigned long long) e.second, e.first / (double) e.second); + } + } + std::vector> v; + { std::lock_guard lk(g_hist_mu); for (auto& kv : g_kernel_counts) v.emplace_back(kv.second, kv.first); } + std::sort(v.begin(), v.end(), [](auto& a, auto& b) { return a.first > b.first; }); + for (size_t i = 0; i < v.size() && i < 14; ++i) + std::fprintf(stderr, " %6llu %s\n", (unsigned long long) v[i].first, v[i].second.c_str()); +} + +extern "C" unsigned long long ds4_dispatch_seq() { return g_disp_count.load(); } + +extern "C" void ds4_gpu_sync_all_streams_weak() __asm("_ds4_gpu_sync_all_streams") __attribute__((weak_import)); +void Context::split_runs() { + // (the streams' pending copies come first, as in a full flush: they are queued on another MTLCommandQueue and would + // otherwise not have run when the commands committed here read what they write) + if (ds4_gpu_sync_all_streams_weak) ds4_gpu_sync_all_streams_weak(); + // Commit the batched dispatches as their own command buffer WITHOUT waiting: the queue runs buffers in order and makes each + // one's writes visible to the next. (Measured: a producer kernel's output was read stale by the next dispatch of the same + // buffer for two kernels in a layer; ending the buffer after them is enough, no wait needed.) + std::lock_guard lk_(impl_->mu); + if (impl_->cb_def == nil) return; + ++g_cb_count; + [impl_->enc_def endEncoding]; + if (impl_->qev == nil) impl_->qev = [impl_->dev newSharedEvent]; + [impl_->cb_def encodeSignalEvent:impl_->qev value:++impl_->qev_v]; + add_busy_handler(impl_->cb_def); + [impl_->cb_def commit]; + impl_->inflight = impl_->cb_def; + impl_->inflight_names = std::move(impl_->def_names); impl_->def_names.clear(); + impl_->cb_def = nil; + impl_->enc_def = nil; + impl_->def_n = 0; +} + +// Waiting for a committed command buffer: spin on its status for a while (the thread wake of waitUntilCompleted costs +// ~1 ms a wait; the decode token waits ~100 times), then block. DS4_NO_SPIN=1: always block. +static void wait_spin(id cb) { + static const bool no_spin = std::getenv("DS4_NO_SPIN") != nullptr; + static const double budget = std::getenv("DS4_SPIN_US") ? std::atof(std::getenv("DS4_SPIN_US")) * 1e-6 : 20e-3; + if (!no_spin) { + const auto t0 = std::chrono::steady_clock::now(); + while (cb.status < MTLCommandBufferStatusCompleted) { + if (std::chrono::duration(std::chrono::steady_clock::now() - t0).count() > budget) break; + } + } + if (cb.status < MTLCommandBufferStatusCompleted) [cb waitUntilCompleted]; +} + +void* Context::mark() { + std::lock_guard lk_(impl_->mu); + if (impl_->cb_def != nil) split_runs(); // (the streams' pending work goes first, inside) + if (impl_->inflight == nil) return nullptr; + return (__bridge_retained void*) impl_->inflight; +} + +void Context::drop_mark(void* token) { + // The token holds a +1 on a committed command buffer; one nobody waits for is never freed. + if (token) { id cb = (__bridge_transfer id) token; cb = nil; } +} + +void Context::wait_mark(void* token) { + id cb = (__bridge_transfer id) token; + const auto tw0 = std::chrono::steady_clock::now(); + wait_spin(cb); + g_wait_ns += (uint64_t) std::chrono::duration_cast(std::chrono::steady_clock::now() - tw0).count(); + ++g_wait_n; + if (cb.status != MTLCommandBufferStatusCompleted) + std::fprintf(stderr, "mtl: deferred command buffer failed: %s\n", [[cb.error localizedDescription] UTF8String]); + if (cb.status != MTLCommandBufferStatusCompleted && std::getenv("DS4_ERRINFO")) std::fprintf(stderr, "mtl: recent: %s\n", recent_str().c_str()), std::fprintf(stderr, "mtl: info: %s\n", [[[cb.error userInfo] description] UTF8String]); +} + +void Context::flush_runs() { + if (std::getenv("DS4_TRACE_DISP") && impl_->cb_def != nil) std::fprintf(stderr, "TRACE flush_runs (%d queued)\n", impl_->def_n); + static const bool cb_log = std::getenv("DS4_CB_LOG") != nullptr; // debugging: one line a command buffer (what it held) + if (cb_log && impl_->cb_def != nil) { + std::string first, last; + if (!impl_->def_names.empty()) { + size_t sp = impl_->def_names.find(' '); + first = impl_->def_names.substr(0, sp); + size_t e = impl_->def_names.find_last_not_of(' '); + size_t b = impl_->def_names.find_last_of(' ', e); + last = impl_->def_names.substr(b == std::string::npos ? 0 : b + 1, e - (b == std::string::npos ? 0 : b)); + } + std::fprintf(stderr, "CB %d disp first=%-24s last=%s\n", impl_->def_n, first.c_str(), last.c_str()); + } + std::lock_guard lk_(impl_->mu); + if (!g_count_reg) { g_count_reg = true; std::atexit(ds4_report_counts); } + if (impl_->cb_def != nil) { + ++g_cb_count; + [impl_->enc_def endEncoding]; + if (impl_->qev == nil) impl_->qev = [impl_->dev newSharedEvent]; + [impl_->cb_def encodeSignalEvent:impl_->qev value:++impl_->qev_v]; + add_busy_handler(impl_->cb_def); + [impl_->cb_def commit]; + impl_->inflight = impl_->cb_def; + impl_->inflight_names = std::move(impl_->def_names); impl_->def_names.clear(); + impl_->cb_def = nil; + impl_->enc_def = nil; + impl_->def_n = 0; + } + if (impl_->inflight != nil) { + if (std::getenv("DS4_TRACE_WAIT")) { void* fr[8]; const int nf = backtrace(fr, 8); std::fprintf(stderr, "WAIT flush_runs:\n"); backtrace_symbols_fd(fr + 1, nf - 1, 2); } + const auto tw0 = std::chrono::steady_clock::now(); + wait_spin(impl_->inflight); + g_wait_ns += (uint64_t) std::chrono::duration_cast(std::chrono::steady_clock::now() - tw0).count(); + ++g_wait_n; + { const double d = impl_->inflight.GPUEndTime - impl_->inflight.GPUStartTime; if (d > 0) g_gpu_ms += d * 1000.0; + if (!g_prof_last.empty()) { std::lock_guard pl(g_prof_mu); auto& e = g_prof[g_prof_last]; e.first += d * 1e6; e.second += 1; g_prof_last.clear(); } } + if (impl_->inflight.status != MTLCommandBufferStatusCompleted) + std::fprintf(stderr, "mtl: deferred command buffer failed: %s\n", [[impl_->inflight.error localizedDescription] UTF8String]); + if (impl_->inflight.status != MTLCommandBufferStatusCompleted && std::getenv("DS4_ERRINFO")) std::fprintf(stderr, "mtl: failed cb held: %s\n", impl_->inflight_names.c_str()); + impl_->inflight = nil; + } + for (void* h : impl_->graveyard) { id x = (__bridge_transfer id)h; x = nil; } + impl_->graveyard.clear(); +#endif +} + +double Context::run(const std::vector& ds, int reps, bool concurrent, bool barrier_between) { +#if defined(STRATA_DS4_RUNTIME) + std::lock_guard lk_(impl_->mu); +#endif + @autoreleasepool { +#if defined(STRATA_DS4_RUNTIME) + if (ds.size() == 1 && reps == 1 && !barrier_between) { + { static const std::string skipk = std::getenv("DS4_SKIP_K") ? std::string(",") + std::getenv("DS4_SKIP_K") + "," : std::string(); // debugging: drop these kernels (timing only) + if (!skipk.empty() && skipk.find("," + ds[0].p->name + ",") != std::string::npos) return 0.0; } + if (impl_->cb_def == nil || impl_->def_n >= 2000000) { + if (impl_->cb_def != nil) split_runs(); // (a buffer committed earlier may still be running: the queue keeps the order) + if (std::getenv("DS4_ERRINFO")) { + MTLCommandBufferDescriptor* cd = [MTLCommandBufferDescriptor new]; + cd.retainedReferences = NO; + cd.errorOptions = MTLCommandBufferErrorOptionEncoderExecutionStatus; + impl_->cb_def = [impl_->q commandBufferWithDescriptor:cd]; + } else + impl_->cb_def = [impl_->q commandBufferWithUnretainedReferences]; + // (the wait on our own event is only for a *different* queue: this queue already runs the previous buffer + // before this one, and the wait costs a full GPU round trip at every buffer - 107 buffers a decode token, + // 0.28 ms each. DS4_QEV_WAIT=1 restores it.) + static const bool qev_wait = std::getenv("DS4_NO_QEV_WAIT") == nullptr; // (measured: off is slower - keep it) + if (qev_wait && impl_->qev != nil && impl_->qev_v > 0) [impl_->cb_def encodeWaitForEvent:impl_->qev value:impl_->qev_v]; + impl_->enc_def = tl_make_encoder(impl_->dev, impl_->cb_def, std::getenv("DS4_CONCURRENT") ? MTLDispatchTypeConcurrent : MTLDispatchTypeSerial); + } + const auto te0 = std::chrono::steady_clock::now(); + encode(impl_->enc_def, ds[0]); + if (g_tl.on && g_tl.att && g_tl.n < g_tl.cap) { g_tl.names[g_tl.n++] = ds[0].p->name; g_tl.att = false; } + g_last_act_ns.store(ka_now_ns(), std::memory_order_relaxed); + g_recent.push_back(ds[0].p->name); if (g_recent.size() > 8) g_recent.pop_front(); + if (std::getenv("DS4_ERRINFO")) impl_->def_names += ds[0].p->name + " "; + g_enc_ns += (uint64_t) std::chrono::duration_cast(std::chrono::steady_clock::now() - te0).count(); + static const bool trace = std::getenv("DS4_TRACE_DISP") != nullptr; + if (trace) { std::lock_guard lk(g_hist_mu); std::fprintf(stderr, "TRACE disp %s\n", ds[0].p->name.c_str()); } + // Every wrapper binds its own MTLBuffer over the same memory, so Metal's hazard tracking cannot see that one + // dispatch reads what the previous wrote: without an explicit barrier the reader saw the stale value (measured: the + // second matvec's activation quantiser read the previous layer's x). + static const bool no_mb = std::getenv("DS4_NO_MBARRIER") != nullptr; + if (!no_mb) [impl_->enc_def memoryBarrierWithScope:MTLBarrierScopeBuffers]; + static const bool enc_each = std::getenv("DS4_ENC_EACH") != nullptr || g_tl.on; // a new encoder per dispatch (debugging, the timeline) + if (enc_each) { + [impl_->enc_def endEncoding]; + impl_->enc_def = tl_make_encoder(impl_->dev, impl_->cb_def, (std::getenv("DS4_CONCURRENT") ? MTLDispatchTypeConcurrent : MTLDispatchTypeSerial)); + } + ++impl_->def_n; + ++g_disp_count; + static const bool prof = std::getenv("DS4_PROF") != nullptr; // every dispatch alone, its GPU time kept per kernel + if (prof) { g_prof_last = ds[0].p->name; if (ds4_gpu_sync_all_streams_weak) ds4_gpu_sync_all_streams_weak(); flush_runs(); return 0.0; } + static const bool each = std::getenv("DS4_FLUSH_EACH") != nullptr; // debugging: every dispatch alone, in order + static const std::string only = std::getenv("DS4_FLUSH_AFTER") ? std::getenv("DS4_FLUSH_AFTER") : ""; // ",name1,name2," + const bool listed = !only.empty() && only.find("," + ds[0].p->name + ",") != std::string::npos; + if (each || listed) { + if (ds4_gpu_sync_all_streams_weak) ds4_gpu_sync_all_streams_weak(); + flush_runs(); + } + return 0.0; + } + flush_runs(); +#endif + id cb = [impl_->q commandBufferWithUnretainedReferences]; + id enc = + [cb computeCommandEncoderWithDispatchType:concurrent ? MTLDispatchTypeConcurrent : MTLDispatchTypeSerial]; + for (int r = 0; r < reps; ++r) + for (const auto& d : ds) { + encode(enc, d); +#if defined(STRATA_DS4_RUNTIME) + ++g_disp_count; +#endif + if (barrier_between) [enc memoryBarrierWithScope:MTLBarrierScopeBuffers]; + } + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + if (cb.status != MTLCommandBufferStatusCompleted) { + std::fprintf(stderr, "mtl: command buffer failed: %s\n", [[cb.error localizedDescription] UTF8String]); + std::exit(5); + } + return cb.GPUEndTime - cb.GPUStartTime; + } +} + +double Context::run_graph(const std::vector& ds) { + @autoreleasepool { + id cb = [impl_->q commandBufferWithUnretainedReferences]; + id enc = [cb computeCommandEncoderWithDispatchType:MTLDispatchTypeConcurrent]; + for (size_t i = 0; i < ds.size(); ++i) { + if (i > 0 && ds[i].barrier) [enc memoryBarrierWithScope:MTLBarrierScopeBuffers]; + encode(enc, ds[i]); + } + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + if (cb.status != MTLCommandBufferStatusCompleted) { + std::fprintf(stderr, "mtl: command buffer failed: %s\n", [[cb.error localizedDescription] UTF8String]); + std::exit(5); + } + return cb.GPUEndTime - cb.GPUStartTime; + } +} + +double Context::run_graph_seg(const std::vector& ds, const std::vector& cuts, void (*at_cut)(void*, int), + void* user) { + @autoreleasepool { + if (impl_->ev == nil) impl_->ev = [impl_->dev newSharedEvent]; + id ev = impl_->ev; + const uint64_t base = impl_->ev_v; + id cb = [impl_->q commandBufferWithUnretainedReferences]; + size_t i = 0; + for (size_t s = 0; s <= cuts.size(); ++s) { + const size_t end = s < cuts.size() ? (size_t) cuts[s] : ds.size(); + id enc = [cb computeCommandEncoderWithDispatchType:MTLDispatchTypeConcurrent]; + const size_t first = i; + for (; i < end; ++i) { + if (i > first && ds[i].barrier) [enc memoryBarrierWithScope:MTLBarrierScopeBuffers]; + encode(enc, ds[i]); + } + [enc endEncoding]; + if (s < cuts.size()) { + [cb encodeSignalEvent:ev value:base + 2 * s + 1]; + [cb encodeWaitForEvent:ev value:base + 2 * s + 2]; + } + } + [cb commit]; + for (size_t s = 0; s < cuts.size(); ++s) { + const uint64_t want = base + 2 * s + 1; + while (ev.signaledValue < want) { + if (cb.status >= MTLCommandBufferStatusCompleted) break; // an error: the wait below reports it + } + at_cut(user, (int) s); + ev.signaledValue = base + 2 * s + 2; + } + impl_->ev_v = base + 2 * cuts.size() + 2; + [cb waitUntilCompleted]; + if (cb.status != MTLCommandBufferStatusCompleted) { + std::fprintf(stderr, "mtl: command buffer failed: %s\n", [[cb.error localizedDescription] UTF8String]); + std::exit(5); + } + return cb.GPUEndTime - cb.GPUStartTime; + } +} + +double Context::run_wall(const std::vector& ds, int reps, double* gpu_seconds) { + double gpu = 0; + const double t0 = now_s(); + for (int r = 0; r < reps; ++r) { + @autoreleasepool { + id cb = [impl_->q commandBufferWithUnretainedReferences]; + id enc = [cb computeCommandEncoder]; + for (const auto& d : ds) encode(enc, d); + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + gpu += cb.GPUEndTime - cb.GPUStartTime; + } + } + if (gpu_seconds) *gpu_seconds = gpu; + return now_s() - t0; +} + +void Context::warm(const std::vector& ds, double seconds) { + const double t0 = now_s(); + while (now_s() - t0 < seconds) run(ds); +} + +double Context::run_icb(const std::vector& ds, int reps, bool barrier_between) { + @autoreleasepool { + size_t maxb = 1; + for (const auto& d : ds) maxb = std::max(maxb, d.args.size()); + MTLIndirectCommandBufferDescriptor* desc = [MTLIndirectCommandBufferDescriptor new]; + desc.commandTypes = MTLIndirectCommandTypeConcurrentDispatch; + desc.inheritBuffers = NO; + desc.inheritPipelineState = NO; + desc.maxKernelBufferBindCount = maxb; + id icb = [impl_->dev newIndirectCommandBufferWithDescriptor:desc + maxCommandCount:ds.size() + options:0]; + NSMutableSet* used = [NSMutableSet new]; + for (size_t i = 0; i < ds.size(); ++i) { + const Dispatch& d = ds[i]; + require_fit(d, "indirect-command dispatch"); + id c = [icb indirectComputeCommandAtIndex:i]; + [c setComputePipelineState:P(d.p)]; + for (size_t a = 0; a < d.args.size(); ++a) { + if (!d.args[a].buf) { + std::fprintf(stderr, "mtl: run_icb needs buffer arguments only\n"); + std::exit(5); + } + size_t o; + id mb = BO(*d.args[a].buf, d.args[a].off, o); + [c setKernelBuffer:mb offset:o atIndex:a]; + [used addObject:mb]; + } + if (d.tg_mem) [c setThreadgroupMemoryLength:(d.tg_mem + 15) / 16 * 16 atIndex:0]; + [c concurrentDispatchThreadgroups:MTLSizeMake(d.grid[0], d.grid[1], d.grid[2]) + threadsPerThreadgroup:MTLSizeMake(d.tg[0], d.tg[1], d.tg[2])]; + if (barrier_between) [c setBarrier]; + } + id cb = [impl_->q commandBufferWithUnretainedReferences]; + id enc = [cb computeCommandEncoderWithDispatchType:MTLDispatchTypeConcurrent]; + for (id b in used) [enc useResource:b usage:MTLResourceUsageRead | MTLResourceUsageWrite]; + for (int r = 0; r < reps; ++r) { + [enc executeCommandsInBuffer:icb withRange:NSMakeRange(0, ds.size())]; + if (barrier_between) [enc memoryBarrierWithScope:MTLBarrierScopeBuffers]; + } + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + if (cb.status != MTLCommandBufferStatusCompleted) { + std::fprintf(stderr, "mtl: icb command buffer failed: %s\n", [[cb.error localizedDescription] UTF8String]); + std::exit(5); + } + return cb.GPUEndTime - cb.GPUStartTime; + } +} + +} // namespace mtl diff --git a/edge0-strata/metal/src/ref_quant.cpp b/edge0-strata/metal/src/ref_quant.cpp new file mode 100644 index 00000000..4d42262a --- /dev/null +++ b/edge0-strata/metal/src/ref_quant.cpp @@ -0,0 +1,517 @@ +// metal/src/ref_quant.cpp - see ref_quant.h. The i-quant loops are transcribed from ggml-quants.c +// (dequantize_row_iq2_xxs / iq2_xs / iq2_s / iq3_xxs / iq3_s / iq4_nl / iq4_xs) at 3cf03257; the K-quants and the +// 32-blocks from Strata's include/strata/artifact/dequant.hpp (themselves transcribed from the same commit). +#include "ref_quant.h" + +#include +#include +#include +#include +#include + +#define GGML_COMMON_DECL_CPP +#define GGML_COMMON_IMPL_CPP +#include "ggml-common.h" +#include "ggml-common-tq.h" // the trellis block structs and the codebook tq_lut_f16 (Gyro) +#include "../../src/kernels/cpu/tq_states.hpp" // the CPU trellis state extraction (the reference decoder) + +namespace ref { + +const char* type_name(int t) { + switch (t) { + case F32: return "F32"; case F16: return "F16"; case BF16: return "BF16"; + case Q4_0: return "Q4_0"; case Q5_0: return "Q5_0"; case Q5_1: return "Q5_1"; case Q8_0: return "Q8_0"; + case Q3_K: return "Q3_K"; case Q4_K: return "Q4_K"; case Q5_K: return "Q5_K"; case Q6_K: return "Q6_K"; + case IQ2_XXS: return "IQ2_XXS"; case IQ2_XS: return "IQ2_XS"; case IQ2_S: return "IQ2_S"; + case IQ3_XXS: return "IQ3_XXS"; case IQ3_S: return "IQ3_S"; case IQ4_NL: return "IQ4_NL"; + case IQ4_XS: return "IQ4_XS"; case Q2_0: return "Q2_0"; case IQ1_S: return "IQ1_S"; case IQ1_M: return "IQ1_M"; + case TQ2_T: return "TQ2_T"; case TQK6: return "TQK6"; case TQK7: return "TQK7"; + } + return "?"; +} + +int block_elems(int t) { + switch (t) { + case F32: case F16: case BF16: return 1; + case Q4_0: case Q5_0: case Q5_1: case Q8_0: case IQ4_NL: return 32; + case Q2_0: return 64; + case TQ2_T: case TQK6: case TQK7: return 128; + default: return 256; + } +} + +int block_bytes(int t) { + switch (t) { + case F32: return 4; case F16: case BF16: return 2; + case Q4_0: return 18; case Q5_0: return 22; case Q5_1: return 24; case Q8_0: return 34; case IQ4_NL: return 18; case Q2_0: return 18; + case Q3_K: return 110; case Q4_K: return 144; case Q5_K: return 176; case Q6_K: return 210; + case IQ2_XXS: return 66; case IQ2_XS: return 74; case IQ2_S: return 82; case IQ3_XXS: return 98; + case IQ3_S: return 110; case IQ4_XS: return 136; case IQ1_S: return 50; case IQ1_M: return 56; + case TQ2_T: return 34; case TQK6: return 26; case TQK7: return 30; + } + std::fprintf(stderr, "ref: unknown type %d\n", t); + std::exit(2); +} + +float f16_to_f32(uint16_t h) { + const uint32_t sign = (uint32_t)(h >> 15) & 1u; + uint32_t exp = (h >> 10) & 0x1Fu, man = h & 0x3FFu, f; + if (exp == 0) { + if (man == 0) + f = sign << 31; + else { + exp = 127 - 15 + 1; + while (!(man & 0x400u)) { man <<= 1; --exp; } + man &= 0x3FFu; + f = (sign << 31) | (exp << 23) | (man << 13); + } + } else if (exp == 31) + f = (sign << 31) | 0x7F800000u | (man << 13); + else + f = (sign << 31) | ((exp - 15 + 127) << 23) | (man << 13); + float out; + std::memcpy(&out, &f, 4); + return out; +} + +uint16_t f32_to_f16(float f) { + _Float16 h = (_Float16)f; + uint16_t r; + std::memcpy(&r, &h, 2); + return r; +} + +float bf16_to_f32(uint16_t h) { + const uint32_t f = (uint32_t)h << 16; + float out; + std::memcpy(&out, &f, 4); + return out; +} + +uint16_t f32_to_bf16(float f) { + uint32_t u; + std::memcpy(&u, &f, 4); + u += 0x7FFFu + ((u >> 16) & 1u); + return (uint16_t)(u >> 16); +} + +static inline uint16_t rd16(const uint8_t* p) { return (uint16_t)(p[0] | (p[1] << 8)); } + +static void get_scale_min_k4(int j, const uint8_t* q, uint8_t& d, uint8_t& m) { + if (j < 4) { + d = q[j] & 63; + m = q[j + 4] & 63; + } else { + d = (q[j + 4] & 0x0F) | ((q[j - 4] >> 6) << 4); + m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); + } +} + +// The trellis types (ggml 144 TQ2_T, 145 TQK6, 146 TQK7): 128 weights = 32 steps of 4; step t's state comes from +// the CPU engine's own extraction (src/kernels/cpu/tq_states.hpp), x = state * 0x9e3779b1, the two codebook points +// lut[x >> 21] and lut[(x >> 10) & 2047] (fp16 pairs) give weights 4t..4t+3, times the fp16 scale. +static void deq_tq(int kb, const uint8_t* b, float* y) { + uint32_t st[32]; + if (kb == 8) strata::kernels::cpu::tq_detail::tq_states<8>(b + 2, st); + else if (kb == 6) strata::kernels::cpu::tq_detail::tq_states<6>(b + 2, st); + else strata::kernels::cpu::tq_detail::tq_states<7>(b + 2, st); + const float d = f16_to_f32(rd16(b)); + for (int t = 0; t < 32; ++t) { + const uint32_t x = st[t] * 0x9e3779b1u; + const uint32_t pa = x >> 21, pb = (x >> 10) & 2047u; + y[4 * t + 0] = f16_to_f32(tq_lut_f16[2 * pa]) * d; + y[4 * t + 1] = f16_to_f32(tq_lut_f16[2 * pa + 1]) * d; + y[4 * t + 2] = f16_to_f32(tq_lut_f16[2 * pb]) * d; + y[4 * t + 3] = f16_to_f32(tq_lut_f16[2 * pb + 1]) * d; + } +} + +static void deq_block(int t, const uint8_t* b, float* y) { + switch (t) { + case TQ2_T: deq_tq(8, b, y); return; + case TQK6: deq_tq(6, b, y); return; + case TQK7: deq_tq(7, b, y); return; + case Q2_0: { + const float d = f16_to_f32(rd16(b)); + for (int j = 0; j < 64; ++j) y[j] = (float)(((b[2 + j / 4] >> ((j % 4) * 2)) & 3) - 1) * d; + return; + } + case Q8_0: { + const float d = f16_to_f32(rd16(b)); + for (int j = 0; j < 32; ++j) y[j] = (float)(int8_t)b[2 + j] * d; + return; + } + case Q4_0: { + const float d = f16_to_f32(rd16(b)); + for (int j = 0; j < 16; ++j) { + y[j] = (float)((b[2 + j] & 15) - 8) * d; + y[j + 16] = (float)((b[2 + j] >> 4) - 8) * d; + } + return; + } + case Q5_0: { + const float d = f16_to_f32(rd16(b)); + uint32_t h; + std::memcpy(&h, b + 2, 4); + const uint8_t* qs = b + 6; + for (int j = 0; j < 16; ++j) { + const int x0 = (qs[j] & 15) | (((h >> j) & 1) << 4), x1 = (qs[j] >> 4) | (((h >> (j + 16)) & 1) << 4); + y[j] = (float)(x0 - 16) * d; + y[j + 16] = (float)(x1 - 16) * d; + } + return; + } + case Q5_1: { // ggml dequantize_row_q5_1: y = x * d + m + const float d = f16_to_f32(rd16(b)), m = f16_to_f32(rd16(b + 2)); + uint32_t h; + std::memcpy(&h, b + 4, 4); + const uint8_t* qs = b + 8; + for (int j = 0; j < 16; ++j) { + const int x0 = (qs[j] & 15) | (((h >> j) & 1) << 4), x1 = (qs[j] >> 4) | (((h >> (j + 16)) & 1) << 4); + y[j] = (float)x0 * d + m; + y[j + 16] = (float)x1 * d + m; + } + return; + } + case IQ4_NL: { + const float d = f16_to_f32(rd16(b)); + for (int j = 0; j < 16; ++j) { + y[j] = d * kvalues_iq4nl[b[2 + j] & 15]; + y[j + 16] = d * kvalues_iq4nl[b[2 + j] >> 4]; + } + return; + } + case IQ4_XS: { + const float d = f16_to_f32(rd16(b)); + const uint16_t sh = rd16(b + 2); + const uint8_t* sl = b + 4; + const uint8_t* qs = b + 8; + for (int ib = 0; ib < 8; ++ib) { + const int ls = ((sl[ib / 2] >> (4 * (ib % 2))) & 15) | (((sh >> (2 * ib)) & 3) << 4); + const float dl = d * (ls - 32); + for (int j = 0; j < 16; ++j) { + y[j] = dl * kvalues_iq4nl[qs[j] & 15]; + y[j + 16] = dl * kvalues_iq4nl[qs[j] >> 4]; + } + y += 32; + qs += 16; + } + return; + } + case Q6_K: { + const uint8_t* ql = b; + const uint8_t* qh = b + 128; + const int8_t* sc = (const int8_t*)(b + 192); + const float d = f16_to_f32(rd16(b + 208)); + for (int n = 0; n < 256; n += 128) { + for (int l = 0; l < 32; ++l) { + const int is = l / 16; + const int q1 = (int)((ql[l] & 15) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int q2 = (int)((ql[l + 32] & 15) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int q3 = (int)((ql[l] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int q4 = (int)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + y[l] = d * sc[is] * q1; + y[l + 32] = d * sc[is + 2] * q2; + y[l + 64] = d * sc[is + 4] * q3; + y[l + 96] = d * sc[is + 6] * q4; + } + y += 128; + ql += 64; + qh += 32; + sc += 8; + } + return; + } + case Q4_K: { + const float d = f16_to_f32(rd16(b)), mn = f16_to_f32(rd16(b + 2)); + const uint8_t* sc = b + 4; + const uint8_t* q = b + 16; + int is = 0; + for (int j = 0; j < 256; j += 64) { + uint8_t s, m; + get_scale_min_k4(is, sc, s, m); + const float d1 = d * s, m1 = mn * m; + get_scale_min_k4(is + 1, sc, s, m); + const float d2 = d * s, m2 = mn * m; + for (int l = 0; l < 32; ++l) *y++ = d1 * (q[l] & 15) - m1; + for (int l = 0; l < 32; ++l) *y++ = d2 * (q[l] >> 4) - m2; + q += 32; + is += 2; + } + return; + } + case Q5_K: { + const float d = f16_to_f32(rd16(b)), mn = f16_to_f32(rd16(b + 2)); + const uint8_t* sc = b + 4; + const uint8_t* qh = b + 16; + const uint8_t* ql = b + 48; + int is = 0; + uint8_t u1 = 1, u2 = 2; + for (int j = 0; j < 256; j += 64) { + uint8_t s, m; + get_scale_min_k4(is, sc, s, m); + const float d1 = d * s, m1 = mn * m; + get_scale_min_k4(is + 1, sc, s, m); + const float d2 = d * s, m2 = mn * m; + for (int l = 0; l < 32; ++l) *y++ = d1 * ((ql[l] & 15) + ((qh[l] & u1) ? 16 : 0)) - m1; + for (int l = 0; l < 32; ++l) *y++ = d2 * ((ql[l] >> 4) + ((qh[l] & u2) ? 16 : 0)) - m2; + ql += 32; + is += 2; + u1 <<= 2; + u2 <<= 2; + } + return; + } + case Q3_K: { + const uint8_t* hm = b; + const uint8_t* q = b + 32; + const float d_all = f16_to_f32(rd16(b + 108)); + const uint32_t kmask1 = 0x03030303u, kmask2 = 0x0f0f0f0fu; + uint32_t aux[4] = {0, 0, 0, 0}; + std::memcpy(aux, b + 96, 12); + const uint32_t tmp = aux[2]; + aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + const int8_t* scales = (const int8_t*)aux; + int is = 0; + uint8_t m = 1; + for (int n = 0; n < 256; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + float dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)); + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) + *y++ = dl * ((int8_t)((q[l + 16] >> shift) & 3) - ((hm[l + 16] & m) ? 0 : 4)); + shift += 2; + m <<= 1; + } + q += 32; + } + return; + } + case IQ1_S: { // ggml dequantize_row_iq1_s: dl = d (2 s + 1), y = dl (grid + delta) + const float d = f16_to_f32(rd16(b)); + const uint8_t* qs = b + 2; + for (int ib = 0; ib < 8; ++ib) { + const uint16_t qh = rd16(b + 34 + 2 * ib); + const float dl = d * (float)(2 * ((qh >> 12) & 7) + 1); + const float delta = qh & 0x8000 ? -IQ1S_DELTA : IQ1S_DELTA; + for (int l = 0; l < 4; ++l) { + const int8_t* grid = (const int8_t*)(iq1s_grid + (qs[l] | (((qh >> 3 * l) & 7) << 8))); + for (int j = 0; j < 8; ++j) y[j] = dl * (grid[j] + delta); + y += 8; + } + qs += 4; + } + return; + } + case IQ1_M: { // ggml dequantize_row_iq1_m: the f16 scale in the top nibbles of the four u16 scales + const uint8_t* qs = b; + const uint8_t* qh = b + 32; + const uint16_t* sc = (const uint16_t*)(b + 48); + uint16_t s16 = (uint16_t)((sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000)); + const float d = f16_to_f32(s16); + for (int ib = 0; ib < 8; ++ib) { + const float dl1 = d * (float)(2 * ((sc[ib / 2] >> (6 * (ib % 2) + 0)) & 7) + 1); + const float dl2 = d * (float)(2 * ((sc[ib / 2] >> (6 * (ib % 2) + 3)) & 7) + 1); + uint16_t idx[4]; + idx[0] = qs[0] | ((qh[0] << 8) & 0x700); + idx[1] = qs[1] | ((qh[0] << 4) & 0x700); + idx[2] = qs[2] | ((qh[1] << 8) & 0x700); + idx[3] = qs[3] | ((qh[1] << 4) & 0x700); + const float delta[4] = {qh[0] & 0x08 ? -IQ1S_DELTA : IQ1S_DELTA, qh[0] & 0x80 ? -IQ1S_DELTA : IQ1S_DELTA, + qh[1] & 0x08 ? -IQ1S_DELTA : IQ1S_DELTA, qh[1] & 0x80 ? -IQ1S_DELTA : IQ1S_DELTA}; + for (int l = 0; l < 4; ++l) { + const int8_t* grid = (const int8_t*)(iq1s_grid + idx[l]); + for (int j = 0; j < 8; ++j) y[j] = (l < 2 ? dl1 : dl2) * (grid[j] + delta[l]); + y += 8; + } + qs += 4; + qh += 2; + } + return; + } + case IQ2_XXS: { + const float d = f16_to_f32(rd16(b)); + const uint8_t* qs = b + 2; + for (int ib32 = 0; ib32 < 8; ++ib32) { + uint32_t aux32[2]; + std::memcpy(aux32, qs + 8 * ib32, 8); + const uint8_t* aux8 = (const uint8_t*)aux32; + const float db = d * (0.5f + (aux32[1] >> 28)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint8_t* grid = (const uint8_t*)(iq2xxs_grid + aux8[l]); + const uint8_t signs = ksigns_iq2xs[(aux32[1] >> 7 * l) & 127]; + for (int j = 0; j < 8; ++j) y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + y += 8; + } + } + return; + } + case IQ2_XS: { + const float d = f16_to_f32(rd16(b)); + const uint8_t* q8 = b + 2; + const uint8_t* scales = b + 66; + for (int ib32 = 0; ib32 < 8; ++ib32) { + const float db0 = d * (0.5f + (scales[ib32] & 15)) * 0.25f; + const float db1 = d * (0.5f + (scales[ib32] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint16_t q = rd16(q8 + 2 * (4 * ib32 + l)); + const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q & 511)); + const uint8_t signs = ksigns_iq2xs[q >> 9]; + const float db = l < 2 ? db0 : db1; + for (int j = 0; j < 8; ++j) y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + y += 8; + } + } + return; + } + case IQ2_S: { + const float d = f16_to_f32(rd16(b)); + const uint8_t* qs = b + 2; + const uint8_t* qh = b + 66; + const uint8_t* scales = b + 74; + const uint8_t* signs = qs + 32; + for (int ib32 = 0; ib32 < 8; ++ib32) { + const float db0 = d * (0.5f + (scales[ib32] & 15)) * 0.25f; + const float db1 = d * (0.5f + (scales[ib32] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const float dl = l < 2 ? db0 : db1; + const uint8_t* grid = (const uint8_t*)(iq2s_grid + (qs[l] | ((qh[ib32] << (8 - 2 * l)) & 0x300))); + for (int j = 0; j < 8; ++j) y[j] = dl * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1.f : 1.f); + y += 8; + } + qs += 4; + signs += 4; + } + return; + } + case IQ3_XXS: { + const float d = f16_to_f32(rd16(b)); + const uint8_t* qs = b + 2; + const uint8_t* sas = b + 2 + 64; + for (int ib32 = 0; ib32 < 8; ++ib32) { + uint32_t aux32; + std::memcpy(&aux32, sas + 4 * ib32, 4); + const float db = d * (0.5f + (aux32 >> 28)) * 0.5f; + for (int l = 0; l < 4; ++l) { + const uint8_t signs = ksigns_iq2xs[(aux32 >> 7 * l) & 127]; + const uint8_t* g1 = (const uint8_t*)(iq3xxs_grid + qs[2 * l + 0]); + const uint8_t* g2 = (const uint8_t*)(iq3xxs_grid + qs[2 * l + 1]); + for (int j = 0; j < 4; ++j) { + y[j + 0] = db * g1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f); + y[j + 4] = db * g2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f); + } + y += 8; + } + qs += 8; + } + return; + } + case IQ3_S: { + const float d = f16_to_f32(rd16(b)); + const uint8_t* qs = b + 2; + const uint8_t* qh = b + 66; + const uint8_t* signs = b + 74; + const uint8_t* scales = b + 106; + for (int ib32 = 0; ib32 < 8; ib32 += 2) { + const float db1 = d * (1 + 2 * (scales[ib32 / 2] & 15)); + const float db2 = d * (1 + 2 * (scales[ib32 / 2] >> 4)); + for (int l = 0; l < 4; ++l) { + const uint8_t* g1 = (const uint8_t*)(iq3s_grid + (qs[2 * l + 0] | ((qh[0] << (8 - 2 * l)) & 256))); + const uint8_t* g2 = (const uint8_t*)(iq3s_grid + (qs[2 * l + 1] | ((qh[0] << (7 - 2 * l)) & 256))); + for (int j = 0; j < 4; ++j) { + y[j + 0] = db1 * g1[j] * (signs[l] & kmask_iq2xs[j + 0] ? -1.f : 1.f); + y[j + 4] = db1 * g2[j] * (signs[l] & kmask_iq2xs[j + 4] ? -1.f : 1.f); + } + y += 8; + } + qs += 8; + signs += 4; + for (int l = 0; l < 4; ++l) { + const uint8_t* g1 = (const uint8_t*)(iq3s_grid + (qs[2 * l + 0] | ((qh[1] << (8 - 2 * l)) & 256))); + const uint8_t* g2 = (const uint8_t*)(iq3s_grid + (qs[2 * l + 1] | ((qh[1] << (7 - 2 * l)) & 256))); + for (int j = 0; j < 4; ++j) { + y[j + 0] = db2 * g1[j] * (signs[l] & kmask_iq2xs[j + 0] ? -1.f : 1.f); + y[j + 4] = db2 * g2[j] * (signs[l] & kmask_iq2xs[j + 4] ? -1.f : 1.f); + } + y += 8; + } + qh += 2; + qs += 8; + signs += 4; + } + return; + } + } + std::fprintf(stderr, "ref: dequant of %s not implemented\n", type_name(t)); + std::exit(2); +} + +void dequant(int t, const void* src, float* dst, size_t n) { + const uint8_t* s = (const uint8_t*)src; + if (t == F32) { std::memcpy(dst, s, n * 4); return; } + if (t == F16) { for (size_t i = 0; i < n; ++i) dst[i] = f16_to_f32(rd16(s + 2 * i)); return; } + if (t == BF16) { for (size_t i = 0; i < n; ++i) dst[i] = bf16_to_f32(rd16(s + 2 * i)); return; } + const int be = block_elems(t), bb = block_bytes(t); + for (size_t i = 0; i < n / be; ++i) deq_block(t, s + i * bb, dst + i * be); +} + +void quantize_q8_0(const float* x, size_t n, uint8_t* out) { + for (size_t b = 0; b < n / 32; ++b) { + float amax = 0.0f; + for (int j = 0; j < 32; ++j) amax = std::max(amax, std::fabs(x[b * 32 + j])); + const float d = amax / 127.0f, id = d != 0.0f ? 1.0f / d : 0.0f; + const uint16_t h = f32_to_f16(d); + uint8_t* o = out + b * 34; + o[0] = (uint8_t)h; + o[1] = (uint8_t)(h >> 8); + for (int j = 0; j < 32; ++j) o[2 + j] = (uint8_t)(int8_t)std::roundf(x[b * 32 + j] * id); + } +} + +static void put16(uint8_t* p, uint16_t v) { p[0] = (uint8_t)v; p[1] = (uint8_t)(v >> 8); } + +void random_blocks(int t, uint8_t* dst, int rows, int cols, uint64_t seed, float scale) { + Rng r(seed); + const size_t n = (size_t)rows * cols; + if (t == F32 || t == F16 || t == BF16) { + for (size_t i = 0; i < n; ++i) { + const float v = r.sym() * scale * 2.0f; + if (t == F32) std::memcpy(dst + 4 * i, &v, 4); + else put16(dst + 2 * i, t == F16 ? f32_to_f16(v) : f32_to_bf16(v)); + } + return; + } + const int be = block_elems(t), bb = block_bytes(t); + const size_t nb = n / be; + for (size_t i = 0; i < nb; ++i) { + uint8_t* b = dst + i * bb; + for (int k = 0; k < bb; ++k) b[k] = (uint8_t)r.next(); + // the scale fields: magnitude so that |w| ~ scale. Typical code magnitudes differ per type. + const float u = 0.5f + r.uni(); + switch (t) { + case Q2_0: put16(b, f32_to_f16(scale * u)); break; + case Q8_0: put16(b, f32_to_f16(scale * u / 64)); break; + case Q4_0: case IQ4_NL: put16(b, f32_to_f16(scale * u / (t == Q4_0 ? 4 : 64))); break; + case Q5_0: put16(b, f32_to_f16(scale * u / 8)); break; + case Q5_1: put16(b, f32_to_f16(scale * u / 16)); put16(b + 2, f32_to_f16(-scale * u)); break; + case IQ4_XS: put16(b, f32_to_f16(scale * u / 64 / 16)); break; + case Q4_K: case Q5_K: put16(b, f32_to_f16(scale * u / 32 / 8)); put16(b + 2, f32_to_f16(scale * u / 32 / 2)); break; + case Q6_K: put16(b + 208, f32_to_f16(scale * u / 32 / 64)); break; + case Q3_K: put16(b + 108, f32_to_f16(scale * u / 32 / 2)); break; + case IQ2_XXS: case IQ2_XS: case IQ2_S: put16(b, f32_to_f16(scale * u / 4)); break; + case IQ1_S: put16(b, f32_to_f16(scale * u / 2)); break; + case IQ1_M: { uint16_t h = f32_to_f16(scale * u / 2); for (int j = 0; j < 4; ++j) { uint16_t w = (uint16_t)(b[48 + 2 * j] | (b[49 + 2 * j] << 8)); w = (uint16_t)((w & 0x0fff) | (((h >> (4 * j)) & 0xf) << 12)); b[48 + 2 * j] = (uint8_t)w; b[49 + 2 * j] = (uint8_t)(w >> 8); } break; } + case IQ3_XXS: put16(b, f32_to_f16(scale * u / 8)); break; + case IQ3_S: put16(b, f32_to_f16(scale * u / 64)); break; + case TQ2_T: case TQK6: case TQK7: put16(b, f32_to_f16(scale * u / 1.5f)); break; // codebook |point| ~ 1 + default: break; + } + } +} + +} // namespace ref diff --git a/edge0-strata/metal/src/ref_quant.h b/edge0-strata/metal/src/ref_quant.h new file mode 100644 index 00000000..5ebe110d --- /dev/null +++ b/edge0-strata/metal/src/ref_quant.h @@ -0,0 +1,47 @@ +// metal/src/ref_quant.h - CPU reference dequantizers for the parity programs, transcribed from ggml's +// dequantize_row_* (ggml-quants.c at the pinned 3cf03257) and Strata's src/artifact/dequant.hpp (Q2_0), plus a +// deterministic generator of random-but-valid blocks for synthetic weights. +#pragma once + +#include +#include +#include + +namespace ref { + +// GGML type ids (stable, as in native_mmvq.hpp) +enum : int { + F32 = 0, F16 = 1, Q4_0 = 2, Q5_0 = 6, Q5_1 = 7, Q8_0 = 8, Q3_K = 11, Q4_K = 12, Q5_K = 13, Q6_K = 14, + IQ2_XXS = 16, IQ2_XS = 17, IQ3_XXS = 18, IQ1_S = 19, IQ1_M = 29, IQ3_S = 21, IQ2_S = 22, IQ4_NL = 20, IQ4_XS = 23, BF16 = 30, Q2_0 = 42, + TQ2_T = 144, TQK6 = 145, TQK7 = 146, // agentionai's trellis types (Gyro): 128 values a block, 34 / 26 / 30 bytes +}; + +const char* type_name(int t); +int block_elems(int t); +int block_bytes(int t); +inline size_t row_bytes(int t, int n) { return (size_t)n / block_elems(t) * block_bytes(t); } + +float f16_to_f32(uint16_t h); +uint16_t f32_to_f16(float f); +float bf16_to_f32(uint16_t h); +uint16_t f32_to_bf16(float f); // round to nearest even + +// Dequantizes `n` elements (a whole number of blocks) of type t. +void dequant(int t, const void* src, float* dst, size_t n); + +// ggml's quantize_row_q8_0_ref: per 32 values d = max|x| / 127 (stored fp16), q = roundf(x * (1 / d)). +void quantize_q8_0(const float* x, size_t n, uint8_t* out); + +// Fills `bytes` with random valid blocks of type t for a matrix of rows x cols: random codes, block scales drawn so +// that the weights have roughly unit-free magnitude `scale` (the d fields are finite FP16). +void random_blocks(int t, uint8_t* dst, int rows, int cols, uint64_t seed, float scale = 0.02f); + +struct Rng { + uint64_t s; + explicit Rng(uint64_t seed) : s(seed * 0x9E3779B97F4A7C15ull + 1) {} + uint64_t next() { s ^= s << 13; s ^= s >> 7; s ^= s << 17; return s; } + float uni() { return (next() >> 40) * (1.0f / 16777216.0f); } // [0,1) + float sym() { return uni() * 2.0f - 1.0f; } // [-1,1) +}; + +} // namespace ref diff --git a/edge0-strata/metal/src/resource_path.h b/edge0-strata/metal/src/resource_path.h new file mode 100644 index 00000000..2e93d57b --- /dev/null +++ b/edge0-strata/metal/src/resource_path.h @@ -0,0 +1,56 @@ +// Where the Metal engine finds the files it reads at run time: the .metal kernel sources it compiles when it starts, +// ggml-common.h, and the default expert-profile / draft-vocabulary tables. +// +// A build used from its own source tree finds them through the paths compiled in (STRATA_*_DEFAULT). A component package +// (strata_control/engine_component.py) is moved away from the tree that built it, so the engine also looks beside its +// own executable first: +// +// /metal/kernels/*.metal the kernel sources (metal/kernels) +// /metal/engine-msl/*.metal the engine's own kernels (metal/src/engine/msl) +// /metal/ggml-common.h +// /metal/data/expert-profile.bin draft_vocab.bin +// +// Order: the environment variable (STRATA_METAL_KERNELS, ...), then that folder when it exists, then the compiled-in +// path. Nothing here touches a build run from its source tree: the folder does not exist there. +#pragma once + +#include +#include + +#include +#include +#include +#include + +namespace strata_metal { + +inline std::string exe_dir() { + char raw[PATH_MAX * 2]; + uint32_t n = sizeof(raw); + if (_NSGetExecutablePath(raw, &n) != 0) return ""; + char real[PATH_MAX]; + if (!realpath(raw, real)) return ""; + std::string s(real); + const size_t slash = s.rfind('/'); + return slash == std::string::npos ? "" : s.substr(0, slash); +} + +inline bool exists(const std::string& path) { + struct stat st; + return stat(path.c_str(), &st) == 0; +} + +// `sub` is relative to /metal/, `compiled` is the path compiled into this build. +inline std::string resource(const char* env, const char* sub, const char* compiled) { + if (env) { + if (const char* e = std::getenv(env)) return e; + } + const std::string dir = exe_dir(); + if (!dir.empty()) { + const std::string beside = dir + "/metal/" + sub; + if (exists(beside)) return beside; + } + return compiled; +} + +} // namespace strata_metal diff --git a/edge0-strata/metal/tests/alu_bench.cpp b/edge0-strata/metal/tests/alu_bench.cpp new file mode 100644 index 00000000..23576071 --- /dev/null +++ b/edge0-strata/metal/tests/alu_bench.cpp @@ -0,0 +1,39 @@ +// metal/tests/alu_bench.cpp - per-operation instruction rates of the GPU (Gops/s, one op = one lane-operation). +#include "mtl_rt.h" + +#include +#include +#include + +int main() { + mtl::Context cx; + const int lib = cx.compile({"alu_bench.metal"}, "#define NO_GGML_COMMON 1\n"); + const uint32_t groups = 2560, tg = 256, n = groups * tg; + mtl::Buffer out = cx.alloc(4 * n), tab = cx.alloc(4096); + for (int i = 0; i < 1024; ++i) tab.as()[i] = i * 7; + struct K { const char* name; double ops_per_iter; } ks[] = { + {"a_fma32", 8}, {"a_fma16", 8}, {"a_fmamix", 8}, {"a_iand", 16}, {"a_shift", 16}, {"a_cvt", 16}, + {"a_tgload", 4}, {"a_sel", 8}, {"a_dot32", 8}, {"a_dot16", 8}, {"a_dotmix", 8}, {"a_fma32x16", 8}, {"a_devload", 4}, {"a_cstload", 4}, {"a_tgload4", 4}, {"a_popc", 4}}; + std::printf("device %s\n", cx.device_name().c_str()); + for (const K& k : ks) { + mtl::Pipeline p = cx.pipeline(lib, k.name); + mtl::Dispatch d; + d.p = &p; + const float kf = 0.999f; + const uint32_t ku = 3; + d.args = {mtl::Arg::B(out), std::string(k.name) == "a_fma32" || std::string(k.name) == "a_fma16" || + std::string(k.name) == "a_fmamix" || std::string(k.name).rfind("a_dot", 0) == 0 || std::string(k.name) == "a_fma32x16" + ? mtl::Arg::V(kf) + : mtl::Arg::V(ku)}; + if (std::string(k.name) == "a_devload" || std::string(k.name) == "a_cstload") d.args.push_back(mtl::Arg::B(tab)); + d.grid[0] = groups; + d.tg[0] = tg; + cx.run({d}); + std::vector t; + for (int i = 0; i < 7; ++i) t.push_back(cx.run({d})); + std::sort(t.begin(), t.end()); + const double ops = (double)n * 1024 * k.ops_per_iter; + std::printf("%-10s %8.1f Gops/s (%.3f ms)\n", k.name, ops / t[3] / 1e9, t[3] * 1e3); + } + return 0; +} diff --git a/edge0-strata/metal/tests/arena_budget_test.cpp b/edge0-strata/metal/tests/arena_budget_test.cpp new file mode 100644 index 00000000..a3be1773 --- /dev/null +++ b/edge0-strata/metal/tests/arena_budget_test.cpp @@ -0,0 +1,135 @@ +// metal/tests/arena_budget_test.cpp - the SSD-streaming arena's size rule (src/engine/arena_budget.h), on the sizes of +// the Macs it was measured on. No GPU: plain arithmetic. Exit 0 = every check passed. +#include "../src/engine/arena_budget.h" + +#include +#include +#include + +using namespace strata::metal_engine; + +static int failures = 0; +static void check(bool ok, const char* what, double got) { + std::printf("%s %-70s %.2f GB\n", ok ? "ok " : "FAIL", what, got / kGB); + failures += !ok; +} + +int main() { + // the M5 Max as it froze: 137.4 GB RAM, 115.4 GB working set, 7.0 GB allocated, 97.7 GB available + ArenaMem m5{137.4 * kGB, 115.4 * kGB, 7.0 * kGB, 97.7 * kGB}; + check(std::fabs(arena_keep(m5) - 0.15 * 137.4 * kGB) < 1e6, "M5 Max: macOS keeps 15% of the RAM (20.6 GB)", arena_keep(m5)); + const double b5 = arena_budget(m5); + check(b5 <= 97.7 * kGB - 20.6 * kGB - 0.5 * kGB + 1e6, "M5 Max: the arena leaves >= 20 GB + the n-gram reserve", b5); + check(b5 > 70 * kGB, "M5 Max: still a large arena (UD-Q5's 98 GB of experts mostly cached)", b5); + // the n-gram table mapped whole (--ple-io ram, 50.7 GB): it comes out of the arena + ArenaMem m5r = m5; + m5r.ple_mapped = 50.7 * kGB; + check(arena_budget(m5r) <= b5 - 48.0 * kGB, "M5 Max, table in RAM: the arena gives up the table's 50.7 GB", arena_budget(m5r)); + // the M4 Mac mini, 32 GB: the 7 GB floor (15% would be 5.16 GB) + ArenaMem m4{34.4 * kGB, 26.8 * kGB, 3.3 * kGB, 17.8 * kGB}; + check(std::fabs(arena_keep(m4) - 7.0 * kGB) < 1e6, "M4 32 GB: macOS keeps the 7 GB floor", arena_keep(m4)); + check(std::fabs(arena_budget(m4) - (17.8 - 7.0 - 0.5) * kGB) < 1e6, "M4 32 GB: available - 7 - 0.5", arena_budget(m4)); + ArenaMem m16{17.2 * kGB, 11 * kGB, 3 * kGB, 9 * kGB}; + check(std::fabs(arena_keep(m16) - 7.0 * kGB) < 1e6, "a 16 GB Mac: the 7 GB floor", arena_keep(m16)); + // the overrides + ArenaMem o = m4; + o.keep_gb = 12.0; + check(std::fabs(arena_budget(o) - (17.8 - 12.0 - 0.5) * kGB) < 1e6, "STRATA_METAL_KEEP_FREE_GB=12", arena_budget(o)); + o.arena_gb = 10.0; + check(std::fabs(arena_budget(o) - 10.0 * kGB) < 1e6, "STRATA_METAL_ARENA_GB=10 wins", arena_budget(o)); + check(std::fabs(arena_shrink(o, 10.0 * kGB, 1.0 * kGB) - 10.0 * kGB) < 1e6, "an explicit arena is not shrunk", 10.0 * kGB); + // never below 1 GB; the GPU working set bounds it too + ArenaMem tiny{16 * kGB, 10 * kGB, 9.5 * kGB, 8 * kGB}; + check(std::fabs(arena_budget(tiny) - 1.0 * kGB) < 1e6, "a 16 GB Mac with nothing left: 1 GB (low hit rate)", arena_budget(tiny)); + // the caps. Anonymous memory (arena, Metal buffers, host memory): 40% of the RAM left on a big Mac; a resident pack: 30%. + check(std::fabs(footprint_cap(m5) - 0.60 * 137.4 * kGB) < 1e6, "M5 Max: the footprint cap is 60% of the RAM (82.4 GB; 40% left)", footprint_cap(m5)); + check(std::fabs(resident_cap(m5) - 0.70 * 137.4 * kGB) < 1e6, "M5 Max: a resident pack may hold 70% (30% left)", resident_cap(m5)); + check(std::fabs(footprint_cap(m4) - 26.8 * kGB) < 1e6, "M4 32 GB: the cap is its 26.8 GB working set (7 GB left, as measured)", footprint_cap(m4)); + check(std::fabs(resident_cap(m4) - 26.8 * kGB) < 1e6, "M4 32 GB: same for a resident pack", resident_cap(m4)); + ArenaMem fc = m5; // the kernel's file-cache floor (37.2 GB here) + 12 GB: never less than that is left + fc.filecache_min = 37.2 * kGB; + check(std::fabs(footprint_cap(fc) - 0.60 * 137.4 * kGB) < 1e6, "M5 Max with its 37.2 GB file-cache floor: 49.2 GB < 55 GB, the share binds", footprint_cap(fc)); + fc.filecache_min = 50.0 * kGB; // a macOS that keeps more: 62 GB left + check(std::fabs(footprint_cap(fc) - (137.4 - 62.0) * kGB) < 1e6, "a larger file-cache floor (50 GB): the cap follows it (75.4 GB)", footprint_cap(fc)); + ArenaMem gpu = m5; + gpu.available = 200 * kGB; + check(std::fabs(arena_budget(gpu) - (0.60 * 137.4 - 7.0 - 0.5) * kGB) < 1e6, "plenty free: the footprint cap bounds it (cap - 7 GB held - 0.5 reserve)", arena_budget(gpu)); + ArenaMem q6 = m5; // UD-Q6_K_XL after a reboot: 7.5 GB allocated, 125 GB "available" (cache); 84 GB ran, 90 GB swapped + q6.allocated = 7.5 * kGB; + q6.available = 125 * kGB; + check(arena_budget(q6) <= 74.5 * kGB && arena_budget(q6) > 74.0 * kGB, "M5 Max, UD-Q6 with 125 GB available: 74.4 GB, not 104", arena_budget(q6)); + // the process holds more than the Metal buffers (host memory): the larger count is the base. UD-Q6 on the M5 held 11.6-15.1 GB + // against 7.5 GB of Metal buffers; the arena gives up the difference, so the total stays at the cap + ArenaMem held = q6; + held.proc_footprint = 15.1 * kGB; + check(std::fabs(arena_budget(held) - (arena_budget(q6) - 7.6 * kGB)) < 1e6, "the process holds 15.1 GB against 7.5 GB of Metal buffers: the arena gives up the 7.6 GB", arena_budget(held)); + FootprintPlan fp = footprint_plan(held, arena_budget(held)); + check(!fp.over() && std::fabs(fp.total - fp.cap) < 1e6, "the planned total (arena + held + n-gram reserve) ends at the cap, not over it", fp.total); + check(137.4 * kGB - fp.total >= 0.40 * 137.4 * kGB - 1e6, "RAM less the planned total is at least 40% of the RAM", 137.4 * kGB - fp.total); + check(std::fabs(fp.total - 82.4 * kGB) < 0.1 * kGB && std::fabs(fp.arena - 66.8 * kGB) < 0.1 * kGB, "= the measured 66.8 GB arena / 82.4 GB total of the clear run", fp.total); + // a footprint smaller than the Metal buffers never makes the arena grow + ArenaMem less = q6; + less.proc_footprint = 1.0 * kGB; + check(std::fabs(arena_budget(less) - arena_budget(q6)) < 1e6, "a smaller process count than the Metal buffers changes nothing", arena_budget(less)); + // little free memory (a 24 GB program already running): the free-memory rule binds below the cap + ArenaMem busy = held; + busy.available = 60 * kGB; + check(arena_budget(busy) < arena_budget(held) && std::fabs(arena_budget(busy) - (60.0 - 20.61 - 0.5) * kGB) < 0.1 * kGB, "a program already holds memory: the free-memory rule binds (38.9 GB)", arena_budget(busy)); + // the explicit arena is kept (never forced) but the plan says it is over the cap + ArenaMem ex = q6; + ex.arena_gb = 90.0; + fp = footprint_plan(ex, arena_budget(ex)); + check(std::fabs(fp.arena - 90.0 * kGB) < 1e6 && fp.over(), "STRATA_METAL_ARENA_GB=90 on the M5: kept, and flagged over the cap", fp.total); + ex.arena_gb = 60.0; + check(!footprint_plan(ex, arena_budget(ex)).over(), "STRATA_METAL_ARENA_GB=60: under the cap", footprint_plan(ex, arena_budget(ex)).total); + ArenaMem sh = q6; // STRATA_METAL_LEAVE_SHARE (for measuring) + sh.leave_share = 0.30; + check(std::fabs(footprint_cap(sh) - 0.70 * 137.4 * kGB) < 1e6, "STRATA_METAL_LEAVE_SHARE=0.30: the old 30% (96.2 GB)", footprint_cap(sh)); + // the working set still binds where it is tighter than the RAM share + ArenaMem m64{68.7 * kGB, 48.0 * kGB, 5.0 * kGB, 55.0 * kGB}; + check(std::fabs(footprint_cap(m64) - 41.2 * kGB) < 0.1 * kGB, "64 GB Mac: 60% of the RAM (41.2 GB), under the 48 GB working set", footprint_cap(m64)); + check(std::fabs(resident_cap(m64) - 48.0 * kGB) < 1e6, "64 GB Mac: a resident pack: the 48 GB working set (30% left would give 48.1)", resident_cap(m64)); + ArenaMem m96{103.1 * kGB, 77.3 * kGB, 6.0 * kGB, 90.0 * kGB}; + check(std::fabs(footprint_cap(m96) - 0.60 * 103.1 * kGB) < 1e6, "96 GB Mac: 60% of the RAM (61.9 GB)", footprint_cap(m96)); + check(std::fabs(resident_cap(m96) - 0.70 * 103.1 * kGB) < 1e6, "96 GB Mac: resident 70% (72.2 GB), under its 77.3 GB working set", resident_cap(m96)); + // a resident pack: its bytes + 12 GiB must fit the resident cap (stream.cpp: wanted) + const double G12 = 12.0 * 1073741824.0; + check(resident_fits(137.4 * kGB, 115.4 * kGB, 54.8 * kGB, G12), "M5 Max: IQ3_S (54.8 GB) stays resident", 54.8 * kGB); + check(resident_fits(137.4 * kGB, 115.4 * kGB, 82.5 * kGB, G12), "M5 Max: UD-Q4_K_XL (82.5 GB) stays resident (measured: wired 91.7 GB, nothing compressed)", 82.5 * kGB); + check(!resident_fits(137.4 * kGB, 115.4 * kGB, 103.9 * kGB, G12), "M5 Max: UD-Q5_K_XL (103.9 GB) streams", 103.9 * kGB); + check(!resident_fits(34.4 * kGB, 26.8 * kGB, 37.6 * kGB, G12), "M4 32 GB: Q2_0 (37.6 GB) streams", 37.6 * kGB); + // memory taken between planning and committing: it shrinks, never grows + const double s1 = arena_shrink(m5, b5, 60.0 * kGB); + check(s1 < b5 && std::fabs(s1 - (60.0 - 20.61 - 0.5) * kGB) < 0.1 * kGB, "less free when committing: it shrinks", s1); + check(std::fabs(arena_shrink(m5, b5, 130.0 * kGB) - b5) < 1e6, "more free when committing: it stays as planned", b5); + // the stepped fill (arena_grow / arena_fill_plan): 1.5 GB steps of 1.75 MB experts on the M4 (7 GB kept + 0.5) + ArenaMem s4 = m4; + s4.available = 17.8 * kGB; + check(arena_grow(s4, 17.8 * kGB, 1.5 * kGB), "a step with room: it fills", 17.8 * kGB); + check(!arena_grow(s4, 8.9 * kGB, 1.5 * kGB), "a step that would leave < 7.5 GB: it stops", 8.9 * kGB); + check(arena_grow(s4, 9.0 * kGB, 1.5 * kGB), "a step that leaves exactly 7.5 GB: it fills", 9.0 * kGB); + ArenaMem pin = s4; + pin.arena_gb = 12.0; + check(arena_grow(pin, 1.0 * kGB, 1.5 * kGB), "STRATA_METAL_ARENA_GB: always fills", 1.0 * kGB); + std::vector cost(6000, 1.75e6); // 10.5 GB of experts planned + // macOS gives exactly what is filled: the plan fills all of it while free stays >= 7.5 GB + const size_t all = arena_fill_plan(s4, cost.data(), cost.size(), 1.5 * kGB, [](double f) { return 18.5 * kGB - f; }); + check(all == cost.size(), "free memory follows the fill: every planned expert is filled", all * 1.75e6); + // another program takes 4 GB while the arena fills (after 3 GB): the fill stops at the floor, a step short + const size_t part = arena_fill_plan(s4, cost.data(), cost.size(), 1.5 * kGB, + [](double f) { return 18.5 * kGB - f - (f >= 3 * kGB ? 4 * kGB : 0); }); + const double left = 18.5 * kGB - part * 1.75e6 - 4 * kGB; + check(part < cost.size() && left >= 7.5 * kGB && left - 1.5 * kGB < 7.5 * kGB + 1.75e6, + "4 GB taken meanwhile: it stops with >= 7.5 GB free, less than a step above", part * 1.75e6); + // no room even for the first step: nothing is filled (the min slots still exist, unfilled) + check(arena_fill_plan(s4, cost.data(), cost.size(), 1.5 * kGB, [](double) { return 8.0 * kGB; }) == 0, + "no room for one step: nothing filled", 0); + // what counts as available: the M4 at the engine's start (16 KiB pages: 6.3 GB free + purgeable, 10.4 GB + // inactive, 14 GB file-backed) - the clean file cache counts; another program's inactive memory does not + check(std::fabs(arena_avail(6.3 * kGB, 0.07 * kGB, 10.4 * kGB, 0.1 * kGB, 14.1 * kGB) - 16.87 * kGB) < 1e6, + "inactive file cache counts", arena_avail(6.3 * kGB, 0.07 * kGB, 10.4 * kGB, 0.1 * kGB, 14.1 * kGB)); + check(std::fabs(arena_avail(6.3 * kGB, 0.07 * kGB, 10.4 * kGB, 0.1 * kGB, 2.0 * kGB) - 8.37 * kGB) < 1e6, + "inactive anonymous memory does not (only 2 GB file-backed)", arena_avail(6.3 * kGB, 0.07 * kGB, 10.4 * kGB, 0.1 * kGB, 2.0 * kGB)); + std::printf("%s\n", failures ? "FAILED" : "all passed"); + return failures ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/coop_layout.cpp b/edge0-strata/metal/tests/coop_layout.cpp new file mode 100644 index 00000000..c296ad5f --- /dev/null +++ b/edge0-strata/metal/tests/coop_layout.cpp @@ -0,0 +1,32 @@ +// metal/tests/coop_layout.cpp - which (index 0, index 1) elements of a single-simdgroup matmul2d's right input, left +// input and destination cooperative tensors each lane holds (mm_bench.metal coop_layout). +// coop_layout [TM_TN_TK] +#include "mtl_rt.h" + +#include +#include + +int main(int argc, char** argv) { + const std::string cfg = argc > 1 ? argv[1] : "16_32_16"; + mtl::Context cx; + const int lib = cx.compile({"mm_bench.metal"}, "#define NO_GGML_COMMON 1\n"); + mtl::Pipeline p = cx.pipeline(lib, "coop_layout_" + cfg); + mtl::Buffer out = cx.alloc(4 * 3 * 32 * 128 * 4); + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(out)}; + d.tg[0] = 32; + cx.run({d}); + const int* o = out.as(); + const char* nm[3] = {"right", "left", "dest"}; + for (int w = 0; w < 3; ++w) { + std::printf("%s (tile %s):\n", nm[w], cfg.c_str()); + for (int l = 0; l < 32; ++l) { + const int* q = o + (w * 32 + l) * 128 * 4; + std::printf(" l%-2d cap %d:", l, q[2]); + for (int i = 0; i < q[2] && i < 128; ++i) std::printf(" (%d,%d)", q[i * 4], q[i * 4 + 1]); + std::printf("\n"); + } + } + return 0; +} diff --git a/edge0-strata/metal/tests/coop_probe.h b/edge0-strata/metal/tests/coop_probe.h new file mode 100644 index 00000000..dc781ff4 --- /dev/null +++ b/edge0-strata/metal/tests/coop_probe.h @@ -0,0 +1,36 @@ +// metal/tests/coop_probe.h - which element layout this GPU gives matmul2d's single-simdgroup right-input cooperative +// tensor: true = runs of 4 consecutive k at one n per lane (the M5), false = anything else (the M4: runs of 2 n at one +// k). The kernels that fill right inputs per lane (mm_gu3, mm_down3, qsa_attn_mm, qsa_flash) take it as function +// constant 7. COOP_K4=0/1 in the environment overrides. +#pragma once +#include "mtl_rt.h" + +#include + +inline bool coop_k4(mtl::Context& cx) { + if (const char* e = std::getenv("COOP_K4")) return std::atoi(e) != 0; + const int lib = cx.compile({"mm_bench.metal"}, "#define NO_GGML_COMMON 1\n"); + for (const char* tile : {"coop_layout_16_64_32", "coop_layout_16_32_32"}) { + mtl::Pipeline p = cx.pipeline(lib, tile); + mtl::Buffer out = cx.alloc(4 * 3 * 32 * 128 * 4); + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(out)}; + d.tg[0] = 32; + cx.run({d}); + // every lane: elements 4j .. 4j + 3 are (k, n), (k + 1, n), (k + 2, n), (k + 3, n), and element i + cap / 2 is + // (k + 16, n) with k % 32 < 16 (mm_gu3's paired-run decode relies on it) + const int* o = out.as(); + for (int l = 0; l < 32; ++l) { + const int* q = o + l * 128 * 4; + const int cap = q[2]; + for (int i = 0; i + 3 < cap; i += 4) + for (int j = 1; j < 4; ++j) + if (q[(i + j) * 4] != q[i * 4] + j || q[(i + j) * 4 + 1] != q[i * 4 + 1]) return false; + for (int i = 0; i < cap / 2; i += 4) + if (q[(i + cap / 2) * 4] != q[i * 4] + 16 || q[(i + cap / 2) * 4 + 1] != q[i * 4 + 1] || (q[i * 4] & 31) >= 16) + return false; + } + } + return true; +} diff --git a/edge0-strata/metal/tests/dense_mm_parity.cpp b/edge0-strata/metal/tests/dense_mm_parity.cpp new file mode 100644 index 00000000..33ce4500 --- /dev/null +++ b/edge0-strata/metal/tests/dense_mm_parity.cpp @@ -0,0 +1,139 @@ +// metal/tests/dense_mm_parity.cpp - dense prompt projections (dense_mm.metal): dequantize to fp16 once, then the +// MPP GEMM. Parity vs a reference with the same fp16 operands (and the exact float one), every dense format and +// awkward row counts; then the speed on the packs' shapes for a 2048-token chunk. +// dense_mm_parity parity +// dense_mm_parity --bench speed +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include + +namespace { +struct DqArgs { uint32_t n_in, n_out, row_bytes, perm; }; +struct GemmArgs { int M, N, K, pad; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +float h16(float v) { return ref::f16_to_f32(ref::f32_to_f16(v)); } +int unit_of(int t) { + switch (t) { + case ref::Q2_0: case ref::IQ3_S: case ref::BF16: case ref::F32: return 16; + default: return 32; + } +} +} // namespace + +int main(int argc, char** argv) { + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); + mtl::Pipeline pact = cx.pipeline(lib, "act_f16"); + auto run_all = [&](int ty, const mtl::Buffer& W, size_t woff, int n_in, int n_out, int T, const mtl::Buffer& x, + const mtl::Buffer& xh, const mtl::Buffer& wh, const mtl::Buffer& y, const std::string& gk, + std::vector& ds, int which) { + static std::vector keep; + mtl::Pipeline* pdq = new mtl::Pipeline(cx.pipeline(lib, "dq_f16_" + std::to_string(ty))); + mtl::Pipeline* pg = new mtl::Pipeline(cx.pipeline(lib, gk)); + keep.push_back(pdq); + keep.push_back(pg); + int tm = 64, tn = 64; + std::sscanf(gk.c_str(), "gemm_f16_%d_%d", &tm, &tn); + if (which & 1) { + const uint32_t n = (uint32_t)T * n_in; + mtl::Dispatch d; d.p = &pact; d.args = {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(n)}; + d.grid[0] = (n + 255) / 256; d.tg[0] = 256; ds.push_back(d); + } + if (which & 2) { + DqArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)ref::row_bytes(ty, n_in)}; + mtl::Dispatch d; d.p = pdq; d.args = {mtl::Arg::B(W, woff), mtl::Arg::B(wh), mtl::Arg::V(a)}; + d.grid[0] = (n_in / unit_of(ty) + 63) / 64; d.grid[1] = n_out; d.tg[0] = 64; ds.push_back(d); + } + if (which & 4) { + GemmArgs g{T, n_out, n_in, 0}; + mtl::Dispatch d; d.p = pg; d.args = {mtl::Arg::B(xh), mtl::Arg::B(wh), mtl::Arg::B(y), mtl::Arg::V(g)}; + d.grid[0] = (n_out + tn - 1) / tn; d.grid[1] = (T + tm - 1) / tm; d.tg[0] = 128; ds.push_back(d); + } + }; + const std::vector types = {ref::F32, ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::IQ4_NL, ref::IQ4_XS, ref::Q2_0, + ref::IQ3_S, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K}; + if (!bench) { + int fails = 0; + for (int ty : types) + for (int n_in : {2560, 640, 6144}) { + if (n_in % 256 && ref::block_elems(ty) == 256) continue; + for (int n_out : {100, 48}) { + const int T = 72; + const size_t rb = ref::row_bytes(ty, n_in); + mtl::Buffer W = cx.alloc(rb * n_out), x = cx.alloc(4 * (size_t)T * n_in), xh = cx.alloc(2 * (size_t)T * n_in), + wh = cx.alloc(2 * (size_t)n_out * n_in), y = cx.alloc(4 * (size_t)T * n_out); + ref::random_blocks(ty, W.as(), n_out, n_in, 5 + ty + n_in, 0.02f); + ref::Rng r(3); + for (size_t i = 0; i < (size_t)T * n_in; ++i) x.as()[i] = r.sym(); + std::vector ds; + run_all(ty, W, 0, n_in, n_out, T, x, xh, wh, y, "gemm_f16_64_64_4", ds, 7); + cx.run(ds); + std::vector wf((size_t)n_in * n_out); + ref::dequant(ty, W.ptr, wf.data(), wf.size()); + double se16 = 0, se = 0, sr = 0; + for (int t = 0; t < T; ++t) + for (int o = 0; o < n_out; ++o) { + double a16 = 0, ae = 0; + for (int k = 0; k < n_in; ++k) { + const float xv = x.as()[(size_t)t * n_in + k], wv = wf[(size_t)o * n_in + k]; + a16 += (double)h16(xv) * h16(wv); + ae += (double)xv * wv; + } + const double g = y.as()[(size_t)t * n_out + o]; + se16 += (g - a16) * (g - a16); + se += (g - ae) * (g - ae); + sr += ae * ae; + } + const double r16 = std::sqrt(se16 / sr), re = std::sqrt(se / sr); + const bool ok = r16 < 1e-5; + fails += !ok; + std::printf("dense mm %-7s %5d -> %4d, T %d: vs fp16 operands %.2e, vs exact %.2e %s\n", + ref::type_name(ty), n_in, n_out, T, r16, re, ok ? "PASS" : "FAIL"); + } + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; + } + struct Shape { const char* name; int type, n_in, n_out; }; + const std::vector shapes = {{"gdn qkv", ref::IQ4_XS, 2560, 10240}, {"gdn qkv", ref::Q3_K, 2560, 10240}, + {"gdn gate", ref::Q4_K, 2560, 6144}, {"gdn out", ref::Q6_K, 6144, 2560}, + {"qsa q", ref::Q2_0, 2560, 12288}, {"hc down", ref::BF16, 10240, 324}, + {"hc up", ref::BF16, 320, 10240}, {"router", ref::BF16, 2560, 512}, + {"alpha", ref::BF16, 2560, 48}}; + const int T = 2048; + for (const Shape& s : shapes) { + if (s.n_in % 128) { std::printf("%-9s %-6s: n_in %d not a multiple of 128 (needs the 32-chunk path)\n", s.name, ref::type_name(s.type), s.n_in); continue; } + const size_t rb = ref::row_bytes(s.type, s.n_in); + mtl::Buffer W = cx.alloc(rb * s.n_out), x = cx.alloc(4 * (size_t)T * s.n_in), xh = cx.alloc(2 * (size_t)T * s.n_in), + wh = cx.alloc(2 * (size_t)s.n_out * s.n_in), y = cx.alloc(4 * (size_t)T * s.n_out); + ref::random_blocks(s.type, W.as(), s.n_out, s.n_in, 3, 0.02f); + auto timeit = [&](int which, const std::string& gk) { + std::vector ds; + for (int i = 0; i < 4; ++i) run_all(s.type, W, 0, s.n_in, s.n_out, T, x, xh, wh, y, gk, ds, which); + cx.warm(ds, 0.2); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds) / 4); + return median(t); + }; + const double tdq = timeit(2, "gemm_f16_64_64_4"), tact = timeit(1, "gemm_f16_64_64_4"); + double best = 1e9; + std::string bk; + for (const char* gk : {"gemm_f16_64_64_4", "gemm_f16_128_64_4", "gemm_f16_64_128_4", "gemm_f16_64_32_4"}) { + const double tg = timeit(4, gk); + if (tg < best) { best = tg; bk = gk; } + } + const double fl = 2.0 * T * s.n_in * s.n_out; + std::printf("%-9s %-6s %5d -> %5d, T %d: dequant %7.1f us, act %6.1f us, GEMM %7.1f us (%5.1f TFLOP/s, %s) -> %.2f us a token\n", + s.name, ref::type_name(s.type), s.n_in, s.n_out, T, tdq * 1e6, tact * 1e6, best * 1e6, fl / best / 1e12, + bk.c_str(), (tdq + best) / T * 1e6); + } + return 0; +} diff --git a/edge0-strata/metal/tests/dsu_job.metal b/edge0-strata/metal/tests/dsu_job.metal new file mode 100644 index 00000000..1a5e028b --- /dev/null +++ b/edge0-strata/metal/tests/dsu_job.metal @@ -0,0 +1,45 @@ +// metal/tests/dsu_job.metal - test-only helpers appended to ds4_iq.metal / moe_mm.metal in dsu_parity.cpp. +// dsu_mkjob writes a Ds4MvJob whose pointers are inside the one buffer it was given (so everything the real kernel +// dereferences is a bound, resident buffer): mem + offsets. +#ifdef DSU_MKJOB +kernel void dsu_mkjob(device uchar* mem [[buffer(0)]], constant ulong* o [[buffer(1)]], constant int* a [[buffer(2)]]) { + // o: job, w, xq, xd, y offsets; a: type, k, n, ncol, ld + device Ds4MvJob& j = *(device Ds4MvJob*) (mem + o[0]); + j.w = mem + o[1]; + j.ld = ulong(a[4]); + j.type = a[0]; + j.k = a[1]; + j.n = a[2]; + j.ncol = a[3]; + j.xq = (device const char*) (mem + o[2]); + j.xd = (device const float*) (mem + o[3]); + j.y = (device float*) (mem + o[4]); + j.ycs = ulong(a[2]); + j.bias = nullptr; + j.alpha = 1.0f; +} +#endif +#ifdef DSU_W8 +// every 4-value run the tile loader reads (k and k + 16 of each 32-group's first half) for T, into half rows +template +static inline void dsu_w8_row(device const uchar* row, int k_total, device half* out, uint gid, threadgroup uint* tab) { + const int runs = k_total / 32 * 4; // per row: 4 runs of 4 values (k = 0,4,8,12 of the group), each giving k and k+16 + if (int(gid) >= runs) return; + const int grp = int(gid) >> 2, r = int(gid) & 3; + const int k = grp * 32 + r * 4; + half4 a, b; + W8::run(row, k, tab, a, b); + out[k] = a.x; out[k + 1] = a.y; out[k + 2] = a.z; out[k + 3] = a.w; + out[k + 16] = b.x; out[k + 17] = b.y; out[k + 18] = b.z; out[k + 19] = b.w; +} +kernel void dsu_w8_iq1s(device const uchar* W [[buffer(0)]], device half* out [[buffer(1)]], constant uint& k [[buffer(2)]], + constant uint& rb [[buffer(3)]], uint2 gid [[thread_position_in_grid]]) { + threadgroup uint tab[1]; + dsu_w8_row(W + ulong(gid.y) * rb, int(k), out + ulong(gid.y) * k, gid.x, tab); +} +kernel void dsu_w8_iq1m(device const uchar* W [[buffer(0)]], device half* out [[buffer(1)]], constant uint& k [[buffer(2)]], + constant uint& rb [[buffer(3)]], uint2 gid [[thread_position_in_grid]]) { + threadgroup uint tab[1]; + dsu_w8_row(W + ulong(gid.y) * rb, int(k), out + ulong(gid.y) * k, gid.x, tab); +} +#endif diff --git a/edge0-strata/metal/tests/dsu_parity.cpp b/edge0-strata/metal/tests/dsu_parity.cpp new file mode 100644 index 00000000..b78d55b5 --- /dev/null +++ b/edge0-strata/metal/tests/dsu_parity.cpp @@ -0,0 +1,191 @@ +// metal/tests/dsu_parity.cpp - the kernels Unsloth's UD-IQ1_S / UD-IQ1_M DeepSeek sets need on the Metal engine, against the CPU +// references in ref_quant.cpp (transcribed from ggml's dequantize_row_*; dsu_ref_vs_ggml checks that transcription against ggml itself): +// 1. ds4_iq1s_mvn1/2/4/8 (IQ1_S expert matvec; IQ1_M and IQ2_XXS as controls) and ds4_iq_mv_k (Q6_K dense matvec): the real engine +// kernels run on int8 activations; relative L2 error against the double-precision dot of the dequantized weights. +// 2. W8 (the prompt tile's weight loader, ds4_mm_gu_19 / ds4_mm_dn_19): bit-exact against the half-rounded reference. +// 3. ds4_dq_q6k (the dense prompt matrices' Q6_K -> half): bit-exact against the half-rounded reference. +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include + +namespace { +int fails = 0; +void report(const char* what, bool ok, double err) { + std::printf("%-62s err %.3e %s\n", what, err, ok ? "PASS" : "FAIL"); + fails += !ok; +} +size_t al(size_t v) { return (v + 255) & ~size_t(255); } + +// activations as the engine's q8 quantiser: per 32 values d = amax / 127, q = round(x / d) +void quant_x(const std::vector& x, int k, int ncol, std::vector& q, std::vector& d) { + q.assign((size_t)ncol * k, 0); + d.assign((size_t)ncol * k / 32, 0.f); + for (int c = 0; c < ncol; ++c) + for (int g = 0; g < k / 32; ++g) { + float am = 0; + for (int i = 0; i < 32; ++i) am = std::fmax(am, std::fabs(x[(size_t)c * k + g * 32 + i])); + const float dd = am / 127.f; + d[(size_t)c * (k / 32) + g] = dd; + for (int i = 0; i < 32; ++i) + q[(size_t)c * k + g * 32 + i] = dd == 0 ? 0 : (int8_t)std::lround(x[(size_t)c * k + g * 32 + i] / dd); + } +} +} // namespace + +int main(int argc, char** argv) { + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + mtl::Context cx; + std::printf("device %s\n", cx.device_name().c_str()); + const int lib = cx.compile({"../src/engine/msl/ds4_iq.metal", "../tests/dsu_job.metal"}, "#define DSU_MKJOB 1\n"); + mtl::Pipeline mk = cx.pipeline(lib, "dsu_mkjob"); + + struct Case { const char* kern; int type; int nc; }; + const std::vector dummy; + const std::vector cases = {{"ds4_iq1s_mvn1", 19, 1}, {"ds4_iq1s_mvn2", 19, 2}, {"ds4_iq1s_mvn4", 19, 4}, {"ds4_iq1s_mvn8", 19, 8}, + {"ds4_iq1m_mvn4", 29, 4}, {"ds4_iq2xxs_mvn4", 16, 4}, {"ds4_iq_mv_k", 14, 8}, {"ds4_iq_mv_k", 14, 1}, + {"ds4_q6k_mvn1", 14, 1}, {"ds4_q6k_mvn2", 14, 2}, {"ds4_q6k_mvn4", 14, 4}, {"ds4_q6k_mvn8", 14, 8}}; + if (bench) { // dense decode matvec speed on a real attention-sized matrix: Q5_K (the 2.00 set's) against Q6_K (the Unsloth sets') + for (const Case& cs : std::vector{{"ds4_q5k_mvn1", 13, 1}, {"ds4_iq_mv_k", 14, 1}, {"ds4_iq1s_mvn1", 19, 1}, {"ds4_iq2xxs_mvn1", 16, 1}, {"ds4_iq_mv_k", 14, 8}, {"ds4_q5k_mvn8", 13, 8}, {"ds4_q6k_mvn1", 14, 1}, {"ds4_q6k_mvn8", 14, 8}}) { + const int k = 4096, n = 16384, nc = cs.nc; + const size_t rb = ref::row_bytes(cs.type, k); + std::vector w(rb * n); + ref::random_blocks(cs.type, w.data(), n, k, 3, 0.05f); + const size_t o_w = al(80), o_xq = al(o_w + w.size()), o_xd = al(o_xq + (size_t)nc * k), o_y = al(o_xd + (size_t)nc * k / 8); + mtl::Buffer mem = cx.alloc(al(o_y + (size_t)nc * n * 4)), offs = cx.alloc(64), args = cx.alloc(64); + std::memcpy(mem.as() + o_w, w.data(), w.size()); + std::memset(mem.as() + o_xq, 3, (size_t)nc * k); + unsigned long ov[5] = {0, o_w, o_xq, o_xd, o_y}; + int av[5] = {cs.type, k, n, nc, (int)rb}; + std::memcpy(offs.ptr, ov, sizeof ov); + std::memcpy(args.ptr, av, sizeof av); + mtl::Pipeline pk = cx.pipeline(lib, cs.kern); + mtl::Dispatch d0; + d0.p = &mk; d0.args = {mtl::Arg::B(mem), mtl::Arg::B(offs), mtl::Arg::B(args)}; d0.grid[0] = 1; d0.tg[0] = 1; + mtl::Dispatch d1; + d1.p = &pk; d1.args = {mtl::Arg::B(mem, 0)}; d1.grid[0] = (n + 7) / 8; d1.grid[1] = 1; d1.tg[0] = 256; + cx.run({d0}); + std::vector ds(20, d1); + cx.warm(ds, 0.3); + double best = 1e9; + for (int i = 0; i < 7; ++i) best = std::fmin(best, cx.run(ds) / 20); + std::printf("bench %-16s %-7s ncol %d: %.1f us %.0f GB/s of weights\n", cs.kern, ref::type_name(cs.type), nc, best * 1e6, (double)w.size() / best / 1e9); + } + return 0; + } + for (const Case& cs : cases) { + for (int k : {512, 4096}) { + const int n = 83; // not a multiple of 8: the last threadgroup's tail rows + const int nc = cs.nc; + const size_t rb = ref::row_bytes(cs.type, k); + std::vector w(rb * n); + ref::random_blocks(cs.type, w.data(), n, k, 11 + cs.type + k, 0.05f); + ref::Rng r(77); + std::vector x((size_t)nc * k); + for (auto& v : x) v = r.sym() * 3.f; + std::vector xq; + std::vector xd; + quant_x(x, k, nc, xq, xd); + const size_t o_job = 0, o_w = al(80), o_xq = al(o_w + w.size()), o_xd = al(o_xq + xq.size()), o_y = al(o_xd + xd.size() * 4); + const size_t total = al(o_y + (size_t)nc * n * 4); + mtl::Buffer mem = cx.alloc(total), offs = cx.alloc(64), args = cx.alloc(64); + std::memcpy(mem.as() + o_w, w.data(), w.size()); + std::memcpy(mem.as() + o_xq, xq.data(), xq.size()); + std::memcpy(mem.as() + o_xd, xd.data(), xd.size() * 4); + unsigned long ov[5] = {o_job, o_w, o_xq, o_xd, o_y}; + int av[5] = {cs.type, k, n, nc, (int)rb}; + std::memcpy(offs.ptr, ov, sizeof ov); + std::memcpy(args.ptr, av, sizeof av); + mtl::Pipeline pk = cx.pipeline(lib, cs.kern); + mtl::Dispatch d0; + d0.p = &mk; d0.args = {mtl::Arg::B(mem), mtl::Arg::B(offs), mtl::Arg::B(args)}; d0.grid[0] = 1; d0.tg[0] = 1; + mtl::Dispatch d1; + d1.p = &pk; d1.args = {mtl::Arg::B(mem, o_job)}; d1.grid[0] = (n + 7) / 8; d1.grid[1] = 1; d1.tg[0] = 256; + cx.run({d0, d1}); + std::vector wf((size_t)n * k); + ref::dequant(cs.type, w.data(), wf.data(), wf.size()); + double se = 0, sr = 0; + const float* y = (const float*)(mem.as() + o_y); + for (int c = 0; c < nc; ++c) + for (int row = 0; row < n; ++row) { + double acc = 0; + for (int i = 0; i < k; ++i) + acc += (double)wf[(size_t)row * k + i] * ((double)xq[(size_t)c * k + i] * xd[(size_t)c * (k / 32) + i / 32]); + const double g = y[(size_t)c * n + row]; + se += (g - acc) * (g - acc); + sr += acc * acc; + } + const double e = std::sqrt(se / sr); + char nm[128]; + std::snprintf(nm, sizeof nm, "%s %s k=%d n=%d ncol=%d", cs.kern, ref::type_name(cs.type), k, n, nc); + report(nm, e < 2e-6, e); + } + } + + // 1b. the engine's own prompt-tile library compiles with the IQ1_S instantiation + { + const int libm = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal", "../src/engine/msl/ds4_mm.metal"}); + for (const char* fn : {"ds4_mm_gu_19", "ds4_mm_dn_19", "ds4_mm_gu_29", "ds4_mm_gu_16"}) { + mtl::Pipeline pp = cx.pipeline(libm, fn); + report(fn, pp.handle != nullptr, 0); + } + } + + // 2. the tile loader + { + const int libw = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal", "../tests/dsu_job.metal"}, "#define DSU_W8 1\n"); + for (int ty : {19, 29}) { + const int k = 2048, n = 37; + const size_t rb = ref::row_bytes(ty, k); + mtl::Buffer W = cx.alloc(rb * n), out = cx.alloc((size_t)n * k * 2); + ref::random_blocks(ty, W.as(), n, k, 5 + ty, 0.05f); + mtl::Pipeline pw = cx.pipeline(libw, ty == 19 ? "dsu_w8_iq1s" : "dsu_w8_iq1m"); + mtl::Dispatch d; + d.p = &pw; d.args = {mtl::Arg::B(W), mtl::Arg::B(out), mtl::Arg::V((uint32_t)k), mtl::Arg::V((uint32_t)rb)}; + d.grid[0] = (k / 32 * 4 + 63) / 64; d.grid[1] = n; d.tg[0] = 64; + cx.run({d}); + std::vector wf((size_t)n * k); + ref::dequant(ty, W.ptr, wf.data(), wf.size()); + size_t bad = 0; + const uint16_t* o = out.as(); + for (size_t i = 0; i < wf.size(); ++i) + if (o[i] != ref::f32_to_f16(wf[i])) ++bad; + char nm[112]; + std::snprintf(nm, sizeof nm, "tile loader W8<%s>, %zu values, mismatching half bits", ref::type_name(ty), wf.size()); + report(nm, bad == 0, (double)bad); + } + } + + // 3. Q6_K -> half + { + const int libd = cx.compile({"../src/engine/msl/ds4_init.metal"}); + mtl::Pipeline pq = cx.pipeline(libd, "ds4_dq_q6k"); + for (int k : {256, 2048}) { + const int n = 45; + const size_t rb = ref::row_bytes(14, k); + mtl::Buffer W = cx.alloc(rb * n), out = cx.alloc((size_t)n * k * 2); + ref::random_blocks(14, W.as(), n, k, 3 + k, 0.05f); + mtl::Dispatch d; + const unsigned long ld = rb; + const uint32_t uk = k, un = n; + d.p = &pq; d.args = {mtl::Arg::B(W), mtl::Arg::V(ld), mtl::Arg::V(uk), mtl::Arg::V(un), mtl::Arg::B(out)}; + d.grid[0] = (uint32_t)(((size_t)n * (k / 256) * 64 + 255) / 256); d.tg[0] = 256; + cx.run({d}); + std::vector wf((size_t)n * k); + ref::dequant(14, W.ptr, wf.data(), wf.size()); + size_t bad = 0; + const uint16_t* o = out.as(); + for (size_t i = 0; i < wf.size(); ++i) + if (o[i] != ref::f32_to_f16(wf[i])) ++bad; + char nm[96]; + std::snprintf(nm, sizeof nm, "ds4_dq_q6k k=%d n=%d, mismatching half bits", k, n); + report(nm, bad == 0, (double)bad); + } + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/dsu_ref_vs_ggml.cpp b/edge0-strata/metal/tests/dsu_ref_vs_ggml.cpp new file mode 100644 index 00000000..d755ab0f --- /dev/null +++ b/edge0-strata/metal/tests/dsu_ref_vs_ggml.cpp @@ -0,0 +1,29 @@ +// metal/tests/dsu_ref_vs_ggml.cpp - ref_quant.cpp's IQ1_S / IQ1_M / Q6_K dequantizers against ggml's own to_float (the pinned llama.cpp +// tree the engine build links), bit for bit. Built only with the engine (needs ggml-base). +#include "ref_quant.h" + +#include "ggml.h" + +#include +#include +#include + +int main() { + int fails = 0; + for (int ty : {ref::IQ1_S, ref::IQ1_M, ref::Q6_K, ref::IQ2_XXS}) { + const int n = 64, k = 2048; + std::vector w(ref::row_bytes(ty, k) * n); + ref::random_blocks(ty, w.data(), n, k, 21 + ty, 0.05f); + std::vector a((size_t)n * k), b((size_t)n * k); + ref::dequant(ty, w.data(), a.data(), a.size()); + const ggml_type_traits* tr = ggml_get_type_traits((ggml_type)ty); + if (!tr || !tr->to_float) { std::printf("%s: ggml has no to_float\n", ref::type_name(ty)); return 2; } + tr->to_float(w.data(), b.data(), (int64_t)b.size()); + const size_t bad = std::memcmp(a.data(), b.data(), a.size() * 4) == 0 ? 0 : 1; + size_t nbad = 0; + for (size_t i = 0; i < a.size(); ++i) nbad += std::memcmp(&a[i], &b[i], 4) != 0; + std::printf("%-8s ref_quant vs ggml to_float: %zu of %zu values differ %s\n", ref::type_name(ty), nbad, a.size(), nbad == 0 ? "PASS" : "FAIL"); + fails += (bad || nbad); + } + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/dsv41/bf16_mv_parity.mm b/edge0-strata/metal/tests/dsv41/bf16_mv_parity.mm new file mode 100644 index 00000000..1045261c --- /dev/null +++ b/edge0-strata/metal/tests/dsv41/bf16_mv_parity.mm @@ -0,0 +1,133 @@ +// bf16mv_probe.mm - the second dsv41 kernel in MSL: the BF16 matrix-vector product (ds4_gguf.cu:1043's mx_mv_kernel / +// bf16_mv_kernel, single-job case). +// +// CUDA's layout: one warp per row, lane i handles k = i, i + 32, ... in BF16 weights times F32 activations, accumulated in +// F32. This compiles the MSL version at run time, runs it on the GPU and compares against the same arithmetic on the CPU. +// clang++ -O2 -fobjc-arc -framework Metal -framework Foundation -o bf16mv_probe bf16mv_probe.mm && ./bf16mv_probe +#import +#import +#include +#include +#include +#include +#include +#include + +static const char* kSrc = R"MSL( +#include +using namespace metal; + +// y[r] = sum_k bf16_to_float(w[r*k + kk]) * x[kk] (one warp a row) +kernel void ds4_bf16_mv(device const ushort* w [[buffer(0)]], device const float* x [[buffer(1)]], + device float* y [[buffer(2)]], constant uint& k [[buffer(3)]], constant uint& n [[buffer(4)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + const uint rows_per_group = 1u; // one warp covers one row + const uint r = gid / 32u * rows_per_group; // the row this warp owns + if (r >= n) return; + float acc = 0.0f; + for (uint kk = lane; kk < k; kk += 32u) { + const uint bits = uint(w[r * k + kk]) << 16; // BF16 -> F32 is the top half + const float wf = as_type(bits); + acc += wf * x[kk]; + } + // warp reduce + acc += simd_shuffle_xor(acc, 16); + acc += simd_shuffle_xor(acc, 8); + acc += simd_shuffle_xor(acc, 4); + acc += simd_shuffle_xor(acc, 2); + acc += simd_shuffle_xor(acc, 1); + if (lane == 0) y[r] = acc; +} +)MSL"; + +static float bf16_to_f32(unsigned short b) { + unsigned bits = (unsigned) b << 16; + float f; + memcpy(&f, &bits, 4); + return f; +} +static unsigned short f32_to_bf16(float f) { + unsigned u; + memcpy(&u, &f, 4); + return (unsigned short) ((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} + +int main() { + @autoreleasepool { + const int n = 2048, k = 5120; + std::vector w((size_t) n * k); + std::vector x(k), y(n); + srand(11); + for (size_t i = 0; i < w.size(); ++i) w[i] = f32_to_bf16(((rand() % 20001) - 10000) / 4096.0f); + for (int i = 0; i < k; ++i) x[i] = ((rand() % 20001) - 10000) / 8192.0f; + + id dev = MTLCreateSystemDefaultDevice(); + NSError* err = nil; + id lib = [dev newLibraryWithSource:[NSString stringWithUTF8String:kSrc] options:nil error:&err]; + if (!lib) { printf("MSL compile failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id pso = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"ds4_bf16_mv"] error:&err]; + if (!pso) { printf("pipeline failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + + id bw = [dev newBufferWithBytes:w.data() length:w.size() * 2 options:MTLResourceStorageModeShared]; + id bx = [dev newBufferWithBytes:x.data() length:k * 4 options:MTLResourceStorageModeShared]; + id by = [dev newBufferWithLength:n * 4 options:MTLResourceStorageModeShared]; + id q = [dev newCommandQueue]; + id cb = [q commandBuffer]; + id ce = [cb computeCommandEncoder]; + [ce setComputePipelineState:pso]; + [ce setBuffer:bw offset:0 atIndex:0]; + [ce setBuffer:bx offset:0 atIndex:1]; + [ce setBuffer:by offset:0 atIndex:2]; + unsigned uk = k, un = n; + [ce setBytes:&uk length:4 atIndex:3]; + [ce setBytes:&un length:4 atIndex:4]; + [ce dispatchThreads:MTLSizeMake((NSUInteger) n * 32, 1, 1) threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + [ce endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + if (cb.error) { printf("dispatch error: %s\n", [[cb.error localizedDescription] UTF8String]); return 1; } + + // throughput: 10 more dispatches, best of them + double best = 1e9; + for (int it = 0; it < 10; ++it) { + id c2 = [q commandBuffer]; + id e2 = [c2 computeCommandEncoder]; + [e2 setComputePipelineState:pso]; + [e2 setBuffer:bw offset:0 atIndex:0]; + [e2 setBuffer:bx offset:0 atIndex:1]; + [e2 setBuffer:by offset:0 atIndex:2]; + [e2 setBytes:&uk length:4 atIndex:3]; + [e2 setBytes:&un length:4 atIndex:4]; + [e2 dispatchThreads:MTLSizeMake((NSUInteger) n * 32, 1, 1) threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + [e2 endEncoding]; + mach_timebase_info_data_t tb; mach_timebase_info(&tb); + const uint64_t t0 = mach_absolute_time(); + [c2 commit]; + [c2 waitUntilCompleted]; + const double dt = (double) (mach_absolute_time() - t0) * tb.numer / tb.denom / 1e9; + if (dt < best) best = dt; + } + { + const double bytes = (double) n * k * 2.0, flops = 2.0 * n * k; + printf("throughput: %.2f GB/s, %.1f GFLOP/s (%.3f ms for 21.0 MB)\n", bytes / best / 1e9, + flops / best / 1e9, best * 1e3); + } + const float* gy = (const float*) [by contents]; + // the standard measure for a dot product: relative L2 over the whole result, not per row (a 5120-term float sum + // cancels, so per-row relative error is ~1e-3 even when the kernels agree to float precision) + double num = 0, den = 0, wrel = 0; + for (int r = 0; r < n; ++r) { + double acc = 0; + for (int kk = 0; kk < k; ++kk) acc += (double) bf16_to_f32(w[(size_t) r * k + kk]) * (double) x[kk]; + const double d = (double) gy[r] - acc; + num += d * d; + den += acc * acc; + const double rel = fabs(d) / (fabs(acc) > 1e-9 ? fabs(acc) : 1.0); + wrel = rel > wrel ? rel : wrel; + } + const double l2 = sqrt(num / den); + printf("rows %d k %d: relative L2 %.3e, worst row %.3e\n", n, k, l2, wrel); + printf("%s\n", l2 < 1e-5 ? "PASS" : "FAIL"); + } + return 0; +} diff --git a/edge0-strata/metal/tests/dsv41/iq2s_t8_parity.mm b/edge0-strata/metal/tests/dsv41/iq2s_t8_parity.mm new file mode 100644 index 00000000..5fd79a6d --- /dev/null +++ b/edge0-strata/metal/tests/dsv41/iq2s_t8_parity.mm @@ -0,0 +1,193 @@ +// iq2s_t8_probe.mm - the IQ2_S expert kernel, tiled over the 8-token window (16 GB of the 2.25bpw set). +// +// Block (block_iq2_s, 82 B): f16 d, 32 index bytes, 32 sign bytes, 8 qh bytes, 8 scale bytes. Per group of 32 weights: +// db0 = d*(0.5 + (scales[ib] & 0xf))*0.25, db1 = d*(0.5 + (scales[ib] >> 4))*0.25 +// for l in 0..3: idx = qs[4*ib + l] | ((qh[ib] << (8 - 2l)) & 0x300); grid = iq2s_grid[idx]; +// sign = signs[4*ib + l]; scale = db[l/2] +// (ggml-quants.c dequantize_row_iq2_s, pinned 3cf03257f; signs live at qs + QK_K/8 in the block.) +// clang++ -O2 -fobjc-arc -framework Metal -framework Foundation -o iq2s_t8_probe iq2s_t8_probe.mm && ./iq2s_t8_probe +#import +#import +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static std::vector table(const std::string& name, size_t expect) { + std::ifstream f("ggml_grids.h"); + std::string all((std::istreambuf_iterator(f)), std::istreambuf_iterator()); + std::regex re("constant \\w+ " + name + "\\[[^\\]]*\\] = \\{([^}]*)\\}"); + std::smatch m; + if (!std::regex_search(all, m, re)) { fprintf(stderr, "no table %s\n", name.c_str()); exit(2); } + std::vector out; + const std::string body = m[1]; + size_t i = 0; + while (i < body.size() && out.size() < expect) { + while (i < body.size() && !isalnum((unsigned char) body[i])) ++i; + size_t j = i; + while (j < body.size() && (isalnum((unsigned char) body[j]) || body[j] == 'x')) ++j; + if (j > i) out.push_back(strtoull(body.substr(i, j - i).c_str(), nullptr, 0)); + i = j; + } + if (out.size() != expect) { fprintf(stderr, "%s: %zu of %zu\n", name.c_str(), out.size(), expect); exit(2); } + return out; +} + +int main() { + @autoreleasepool { + const int n = 65536, k = 5120, T = 8; + const int nb = k / 256; + const size_t bs = 82; + auto grid = table("iq2s_grid", 1024); + + std::vector w((size_t) n * nb * bs); + srand(13); + for (size_t i = 0; i < w.size(); ++i) w[i] = (uint8_t) (rand() & 0xff); + for (int r = 0; r < n; ++r) + for (int b = 0; b < nb; ++b) { + _Float16 d = (_Float16) (((rand() % 900) + 100) / 8000.0f); + memcpy(&w[((size_t) r * nb + b) * bs], &d, 2); + } + std::vector x((size_t) T * k), y((size_t) T * n); + for (size_t i = 0; i < x.size(); ++i) x[i] = ((rand() % 20001) - 10000) / 8192.0f; + + std::string grids; + { + std::ifstream f("ggml_grids.h"); + std::string line; + while (std::getline(f, line)) { + if (line.rfind("#pragma once", 0) == 0 || line.rfind("//", 0) == 0 || line.empty()) continue; + grids += line + "\n"; + } + } + std::string src = "#include \nusing namespace metal;\n" + grids + R"MSL( +kernel void iq2s_t8(device const uchar* w [[buffer(0)]], device const float* x [[buffer(1)]], + device float* y [[buffer(2)]], constant uint& k [[buffer(3)]], constant uint& n [[buffer(4)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + const uint r = gid / 32u; + if (r >= n) return; + const uint nb = k / 256u; + const device uchar* base = w + (size_t) r * nb * 82u; + float acc[8]; + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (uint b = lane; b < nb; b += 32u) { + const device uchar* blk = base + (size_t) b * 82u; + const float d = float(as_type(*(const device ushort*) blk)); + const device uchar* qs = blk + 2; // 32 index bytes + const device uchar* sg = qs + 32; // 32 sign bytes + const device uchar* qh = blk + 2 + 64; // 8 qh bytes + const device uchar* sc = blk + 2 + 72; // 8 scale bytes + for (int ib = 0; ib < 8; ++ib) { + const float db0 = d * (0.5f + float(sc[ib] & 0xf)) * 0.25f; + const float db1 = d * (0.5f + float(sc[ib] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint idx = uint(qs[4 * ib + l]) | ((uint(qh[ib]) << (8 - 2 * l)) & 0x300u); + const ulong g = iq2s_grid[idx & 1023u]; + const uint signs = uint(sg[4 * ib + l]); + const float db = (l < 2) ? db0 : db1; + const uint xoff = (uint) ib * 32u + (uint) l * 8u; + float wv[8]; + for (int j = 0; j < 8; ++j) { + wv[j] = db * float((g >> (8 * j)) & 0xff) * ((signs & kmask_iq2xs[j]) ? -1.0f : 1.0f); + } + for (uint t = 0; t < 8u; ++t) { + const device float* xt = x + (size_t) t * k + xoff; + float s = 0.0f; + for (int j = 0; j < 8; ++j) s += wv[j] * xt[j]; + acc[t] += s; + } + } + } + } + for (uint t = 0; t < 8u; ++t) { + float a = acc[t]; + a += simd_shuffle_xor(a, 16); + a += simd_shuffle_xor(a, 8); + a += simd_shuffle_xor(a, 4); + a += simd_shuffle_xor(a, 2); + a += simd_shuffle_xor(a, 1); + if (lane == 0) y[(size_t) t * n + r] = a; + } +} +)MSL"; + + id dev = MTLCreateSystemDefaultDevice(); + NSError* err = nil; + id lib = [dev newLibraryWithSource:[NSString stringWithUTF8String:src.c_str()] options:nil error:&err]; + if (!lib) { printf("MSL compile failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id pso = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"iq2s_t8"] error:&err]; + if (!pso) { printf("pipeline failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id bw = [dev newBufferWithBytes:w.data() length:w.size() options:MTLResourceStorageModeShared]; + id bx = [dev newBufferWithBytes:x.data() length:x.size() * 4 options:MTLResourceStorageModeShared]; + id by = [dev newBufferWithLength:y.size() * 4 options:MTLResourceStorageModeShared]; + id q = [dev newCommandQueue]; + unsigned uk = k, un = n; + double best = 1e9; + mach_timebase_info_data_t tb; mach_timebase_info(&tb); + for (int it = 0; it < 5; ++it) { + id cb = [q commandBuffer]; + id ce = [cb computeCommandEncoder]; + [ce setComputePipelineState:pso]; + [ce setBuffer:bw offset:0 atIndex:0]; + [ce setBuffer:bx offset:0 atIndex:1]; + [ce setBuffer:by offset:0 atIndex:2]; + [ce setBytes:&uk length:4 atIndex:3]; + [ce setBytes:&un length:4 atIndex:4]; + [ce dispatchThreads:MTLSizeMake((NSUInteger) n * 32, 1, 1) threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + [ce endEncoding]; + const uint64_t t0 = mach_absolute_time(); + [cb commit]; + [cb waitUntilCompleted]; + const double dt = (double) (mach_absolute_time() - t0) * tb.numer / tb.denom / 1e9; + if (dt < best) best = dt; + } + printf("IQ2_S window of 8: %.2f GB/s effective (%.2f raw), %.3f ms for %.1f MB\n", 8.0 * w.size() / best / 1e9, + (double) w.size() / best / 1e9, best * 1e3, (double) w.size() / 1e6); + + const float* gy = (const float*) [by contents]; + const uint8_t* kmask = (const uint8_t*) "\x01\x02\x04\x08\x10\x20\x40\x80"; + double num = 0, den = 0; + for (int r = 0; r < n; r += 61) { + for (int t = 0; t < T; ++t) { + double acc = 0; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = &w[((size_t) r * nb + b) * bs]; + _Float16 dh; memcpy(&dh, blk, 2); + const double d = (double) (float) dh; + const uint8_t* qs = blk + 2; + const uint8_t* sg = qs + 32; + const uint8_t* qh = blk + 2 + 64; + const uint8_t* sc = blk + 2 + 72; + for (int ib = 0; ib < 8; ++ib) { + const double db0 = d * (0.5 + (double) (sc[ib] & 0xf)) * 0.25; + const double db1 = d * (0.5 + (double) (sc[ib] >> 4)) * 0.25; + for (int l = 0; l < 4; ++l) { + const uint32_t idx = (uint32_t) qs[4 * ib + l] | (((uint32_t) qh[ib] << (8 - 2 * l)) & 0x300u); + const uint64_t g = grid[idx & 1023u]; + const uint32_t signs = sg[4 * ib + l]; + const double db = (l < 2) ? db0 : db1; + const int xoff = ib * 32 + l * 8; + for (int j = 0; j < 8; ++j) { + const double wv = db * (double) ((g >> (8 * j)) & 0xff) * ((signs & kmask[j]) ? -1.0 : 1.0); + acc += wv * (double) x[(size_t) t * k + xoff + j]; + } + } + } + } + const double diff = (double) gy[(size_t) t * n + r] - acc; + num += diff * diff; + den += acc * acc; + } + } + const double l2 = sqrt(num / den); + printf("relative L2 %.3e (sampled rows)\n", l2); + printf("%s\n", l2 < 1e-5 ? "PASS" : "FAIL"); + } + return 0; +} diff --git a/edge0-strata/metal/tests/dsv41/iq2xs_t8_parity.mm b/edge0-strata/metal/tests/dsv41/iq2xs_t8_parity.mm new file mode 100644 index 00000000..8f5d1a81 --- /dev/null +++ b/edge0-strata/metal/tests/dsv41/iq2xs_t8_parity.mm @@ -0,0 +1,190 @@ +// iq2xs_t8_probe.mm - the IQ2_XS expert kernel, tiled over the engine's 8-token window (24 GB of the 2.25bpw set). +// +// Block (block_iq2_xs, 74 B): f16 d, 32 uint16 qs, 8 bytes of scales. Per group of 32 weights: +// db0 = d*(0.5 + (scales[ib] & 0xf))*0.25, db1 = d*(0.5 + (scales[ib] >> 4))*0.25 +// for l in 0..3: grid = iq2xs_grid[qs[4ib+l] & 511], signs = ksigns_iq2xs[qs[4ib+l] >> 9], scale = db[l/2] +// (ggml-quants.c dequantize_row_iq2_xs, pinned 3cf03257f.) Decoded once per window, applied to 8 activations. +// clang++ -O2 -fobjc-arc -framework Metal -framework Foundation -o iq2xs_t8_probe iq2xs_t8_probe.mm && ./iq2xs_t8_probe +#import +#import +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static std::vector table(const std::string& name, size_t expect) { + std::ifstream f("ggml_grids.h"); + std::string all((std::istreambuf_iterator(f)), std::istreambuf_iterator()); + std::regex re("constant \\w+ " + name + "\\[[^\\]]*\\] = \\{([^}]*)\\}"); + std::smatch m; + if (!std::regex_search(all, m, re)) { fprintf(stderr, "no table %s\n", name.c_str()); exit(2); } + std::vector out; + const std::string body = m[1]; + size_t i = 0; + while (i < body.size() && out.size() < expect) { + while (i < body.size() && !isalnum((unsigned char) body[i])) ++i; + size_t j = i; + while (j < body.size() && (isalnum((unsigned char) body[j]) || body[j] == 'x')) ++j; + if (j > i) out.push_back(strtoull(body.substr(i, j - i).c_str(), nullptr, 0)); + i = j; + } + if (out.size() != expect) { fprintf(stderr, "%s: %zu of %zu\n", name.c_str(), out.size(), expect); exit(2); } + return out; +} + +int main() { + @autoreleasepool { + const int n = 65536, k = 5120, T = 8; + const int nb = k / 256; + const size_t bs = 74; + auto grid = table("iq2xs_grid", 512); + auto signs_t = table("ksigns_iq2xs", 128); + + std::vector w((size_t) n * nb * bs); + srand(9); + for (size_t i = 0; i < w.size(); ++i) w[i] = (uint8_t) (rand() & 0xff); + for (int r = 0; r < n; ++r) + for (int b = 0; b < nb; ++b) { + _Float16 d = (_Float16) (((rand() % 900) + 100) / 8000.0f); + memcpy(&w[((size_t) r * nb + b) * bs], &d, 2); + } + std::vector x((size_t) T * k), y((size_t) T * n); + for (size_t i = 0; i < x.size(); ++i) x[i] = ((rand() % 20001) - 10000) / 8192.0f; + + std::string grids; + { + std::ifstream f("ggml_grids.h"); + std::string line; + while (std::getline(f, line)) { + if (line.rfind("#pragma once", 0) == 0 || line.rfind("//", 0) == 0 || line.empty()) continue; + grids += line + "\n"; + } + } + std::string src = "#include \nusing namespace metal;\n" + grids + R"MSL( +kernel void iq2xs_t8(device const uchar* w [[buffer(0)]], device const float* x [[buffer(1)]], + device float* y [[buffer(2)]], constant uint& k [[buffer(3)]], constant uint& n [[buffer(4)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + const uint r = gid / 32u; + if (r >= n) return; + const uint nb = k / 256u; + const device uchar* base = w + (size_t) r * nb * 74u; + float acc[8]; + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (uint b = lane; b < nb; b += 32u) { + const device uchar* blk = base + (size_t) b * 74u; + const float d = float(as_type(*(const device ushort*) blk)); + const device ushort* qs = (const device ushort*) (blk + 2); + const device uchar* sc = blk + 2 + 64; + for (int ib = 0; ib < 8; ++ib) { + const float db0 = d * (0.5f + float(sc[ib] & 0xf)) * 0.25f; + const float db1 = d * (0.5f + float(sc[ib] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint q = uint(qs[4 * ib + l]); + const ulong g = iq2xs_grid[q & 511u]; + const uint sg = uint(ksigns_iq2xs[q >> 9]); + const float db = (l < 2) ? db0 : db1; + const uint xoff = (uint) ib * 32u + (uint) l * 8u; + float wv[8]; + for (int j = 0; j < 8; ++j) { + wv[j] = db * float((g >> (8 * j)) & 0xff) * ((sg & kmask_iq2xs[j]) ? -1.0f : 1.0f); + } + for (uint t = 0; t < 8u; ++t) { + const device float* xt = x + (size_t) t * k + xoff; + float s = 0.0f; + for (int j = 0; j < 8; ++j) s += wv[j] * xt[j]; + acc[t] += s; + } + } + } + } + for (uint t = 0; t < 8u; ++t) { + float a = acc[t]; + a += simd_shuffle_xor(a, 16); + a += simd_shuffle_xor(a, 8); + a += simd_shuffle_xor(a, 4); + a += simd_shuffle_xor(a, 2); + a += simd_shuffle_xor(a, 1); + if (lane == 0) y[(size_t) t * n + r] = a; + } +} +)MSL"; + + id dev = MTLCreateSystemDefaultDevice(); + NSError* err = nil; + id lib = [dev newLibraryWithSource:[NSString stringWithUTF8String:src.c_str()] options:nil error:&err]; + if (!lib) { printf("MSL compile failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id pso = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"iq2xs_t8"] error:&err]; + if (!pso) { printf("pipeline failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id bw = [dev newBufferWithBytes:w.data() length:w.size() options:MTLResourceStorageModeShared]; + id bx = [dev newBufferWithBytes:x.data() length:x.size() * 4 options:MTLResourceStorageModeShared]; + id by = [dev newBufferWithLength:y.size() * 4 options:MTLResourceStorageModeShared]; + id q = [dev newCommandQueue]; + unsigned uk = k, un = n; + double best = 1e9; + mach_timebase_info_data_t tb; mach_timebase_info(&tb); + for (int it = 0; it < 5; ++it) { + id cb = [q commandBuffer]; + id ce = [cb computeCommandEncoder]; + [ce setComputePipelineState:pso]; + [ce setBuffer:bw offset:0 atIndex:0]; + [ce setBuffer:bx offset:0 atIndex:1]; + [ce setBuffer:by offset:0 atIndex:2]; + [ce setBytes:&uk length:4 atIndex:3]; + [ce setBytes:&un length:4 atIndex:4]; + [ce dispatchThreads:MTLSizeMake((NSUInteger) n * 32, 1, 1) threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + [ce endEncoding]; + const uint64_t t0 = mach_absolute_time(); + [cb commit]; + [cb waitUntilCompleted]; + const double dt = (double) (mach_absolute_time() - t0) * tb.numer / tb.denom / 1e9; + if (dt < best) best = dt; + } + printf("IQ2_XS window of 8: %.2f GB/s effective (%.2f raw), %.3f ms for %.1f MB\n", 8.0 * w.size() / best / 1e9, + (double) w.size() / best / 1e9, best * 1e3, (double) w.size() / 1e6); + + const float* gy = (const float*) [by contents]; + const uint8_t* kmask = (const uint8_t*) "\x01\x02\x04\x08\x10\x20\x40\x80"; + double num = 0, den = 0; + for (int r = 0; r < n; r += 61) { + for (int t = 0; t < T; ++t) { + double acc = 0; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = &w[((size_t) r * nb + b) * bs]; + _Float16 dh; memcpy(&dh, blk, 2); + const double d = (double) (float) dh; + const uint8_t* sc = blk + 2 + 64; + const uint8_t* qb = blk + 2; + for (int ib = 0; ib < 8; ++ib) { + const double db0 = d * (0.5 + (double) (sc[ib] & 0xf)) * 0.25; + const double db1 = d * (0.5 + (double) (sc[ib] >> 4)) * 0.25; + for (int l = 0; l < 4; ++l) { + uint16_t q16; + memcpy(&q16, qb + 2 * (4 * ib + l), 2); + const uint64_t g = grid[q16 & 511u]; + const uint32_t sg = (uint32_t) signs_t[(q16 >> 9) & 127u]; + const double db = (l < 2) ? db0 : db1; + const int xoff = ib * 32 + l * 8; + for (int j = 0; j < 8; ++j) { + const double wv = db * (double) ((g >> (8 * j)) & 0xff) * ((sg & kmask[j]) ? -1.0 : 1.0); + acc += wv * (double) x[(size_t) t * k + xoff + j]; + } + } + } + } + const double diff = (double) gy[(size_t) t * n + r] - acc; + num += diff * diff; + den += acc * acc; + } + } + const double l2 = sqrt(num / den); + printf("relative L2 %.3e (sampled rows)\n", l2); + printf("%s\n", l2 < 1e-5 ? "PASS" : "FAIL"); + } + return 0; +} diff --git a/edge0-strata/metal/tests/dsv41/iq2xxs_gemv_parity.mm b/edge0-strata/metal/tests/dsv41/iq2xxs_gemv_parity.mm new file mode 100644 index 00000000..19956302 --- /dev/null +++ b/edge0-strata/metal/tests/dsv41/iq2xxs_gemv_parity.mm @@ -0,0 +1,185 @@ +// iq2xxs_probe.mm - the IQ2_XXS expert GEMV in MSL, the 2.25bpw set's dominant expert format (49 GB of the 149 GB). +// +// Decode (ggml-quants.c dequantize_row_iq2_xxs, pinned 3cf03257f): a block is 66 bytes - an f16 d and 32 bytes of qs. Eight +// groups of 32 weights: aux32[0..1] are the group's 8 bytes; db = d * (0.5 + (aux32[1] >> 28)) * 0.25; each of the group's four +// sub-blocks takes grid = iq2xxs_grid[aux8[l]] (8 bytes) and signs = ksigns_iq2xs[(aux32[1] >> 7l) & 127], giving +// y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1). +// +// The grid tables are read from the generated MSL header so the CPU reference uses exactly the same numbers as the GPU. +// clang++ -O2 -fobjc-arc -framework Metal -framework Foundation -o iq2xxs_probe iq2xxs_probe.mm && ./iq2xxs_probe +#import +#import +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static std::vector load_msl_table(const std::string& path, const std::string& name, size_t expect) { + std::ifstream f(path); + std::string all((std::istreambuf_iterator(f)), std::istreambuf_iterator()); + std::regex re("constant \\w+ " + name + "\\[[^\\]]*\\] = \\{([^}]*)\\}"); + std::smatch m; + std::vector out; + if (!std::regex_search(all, m, re)) { fprintf(stderr, "table %s not found\n", name.c_str()); exit(2); } + std::string body = m[1]; + size_t i = 0; + while (i < body.size()) { + while (i < body.size() && !isdigit((unsigned char) body[i])) ++i; + if (i >= body.size()) break; + size_t j = i; + while (j < body.size() && (isdigit((unsigned char) body[j]) || body[j] == 'x' || isxdigit((unsigned char) body[j]))) ++j; + out.push_back(strtoull(body.substr(i, j - i).c_str(), nullptr, 0)); + i = j; + if (out.size() > expect) break; + } + if (out.size() != expect) { fprintf(stderr, "table %s: %zu of %zu\n", name.c_str(), out.size(), expect); exit(2); } + return out; +} + +int main() { + @autoreleasepool { + const int n = 65536, k = 5120; // rows and k + const int nb = k / 256; // IQ2_XXS blocks a row + const size_t bs = 66; // bytes a block + auto grid = load_msl_table("ggml_grids.h", "iq2xxs_grid", 256); + auto signs_t = load_msl_table("ggml_grids.h", "ksigns_iq2xs", 128); + + // a synthetic pack of blocks (the decoder does not care where the bytes come from) + std::vector w((size_t) n * nb * bs); + srand(3); + for (size_t i = 0; i < w.size(); ++i) w[i] = (uint8_t) (rand() & 0xff); + for (int r = 0; r < n; ++r) + for (int b = 0; b < nb; ++b) { + // a sane f16 scale in the block's first two bytes + _Float16 d = (_Float16) (((rand() % 900) + 100) / 8000.0f); + memcpy(&w[((size_t) r * nb + b) * bs], &d, 2); + } + std::vector x(k), y(n); + for (int i = 0; i < k; ++i) x[i] = ((rand() % 20001) - 10000) / 8192.0f; + + // pull the grid table into the MSL by textual include at compile time + std::string grids; + { + std::ifstream f("ggml_grids.h"); + std::string line; + while (std::getline(f, line)) { + if (line.rfind("#pragma once", 0) == 0 || line.rfind("//", 0) == 0 || line.empty()) continue; + grids += line + "\n"; + } + } + std::string src = "#include \nusing namespace metal;\n" + grids + R"MSL( +kernel void iq2xxs_gemv(device const uchar* w [[buffer(0)]], device const float* x [[buffer(1)]], + device float* y [[buffer(2)]], constant uint& k [[buffer(3)]], constant uint& n [[buffer(4)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + const uint r = gid / 32u; + if (r >= n) return; + const uint nb = k / 256u; + const device uchar* base = w + (size_t) r * nb * 66u; + float acc = 0.0f; + for (uint b = lane; b < nb; b += 32u) { + const device uchar* blk = base + (size_t) b * 66u; + const float d = float(as_type(*(const device ushort*) blk)); + for (int ib = 0; ib < 8; ++ib) { + // groups advance 4 bytes and read 8, so an 8-byte vector load is misaligned on odd groups: read bytes + const uint aux0 = uint(blk[2 + 4 * ib]) | (uint(blk[3 + 4 * ib]) << 8) | (uint(blk[4 + 4 * ib]) << 16) | (uint(blk[5 + 4 * ib]) << 24); + const uint aux1 = uint(blk[6 + 4 * ib]) | (uint(blk[7 + 4 * ib]) << 8) | (uint(blk[8 + 4 * ib]) << 16) | (uint(blk[9 + 4 * ib]) << 24); + const float db = d * (0.5f + float(aux1 >> 28)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const ulong g = iq2xxs_grid[(aux0 >> (8 * l)) & 0xff]; + const uint sg = uint(ksigns_iq2xs[(aux1 >> (7 * l)) & 127]); + const uint xoff = (uint) ib * 32u + (uint) l * 8u; + for (int j = 0; j < 8; ++j) { + const float wv = db * float((g >> (8 * j)) & 0xff) * ((sg & kmask_iq2xs[j]) ? -1.0f : 1.0f); + acc += wv * x[xoff + j]; + } + } + } + } + acc += simd_shuffle_xor(acc, 16); + acc += simd_shuffle_xor(acc, 8); + acc += simd_shuffle_xor(acc, 4); + acc += simd_shuffle_xor(acc, 2); + acc += simd_shuffle_xor(acc, 1); + if (lane == 0) y[r] = acc; +} +)MSL"; + const uint8_t kmask[8] = {1, 2, 4, 8, 16, 32, 64, 128}; + // kmask_iq2xs comes from ggml_grids.h + std::string src_with_km = src; + + id dev = MTLCreateSystemDefaultDevice(); + NSError* err = nil; + id lib = [dev newLibraryWithSource:[NSString stringWithUTF8String:src_with_km.c_str()] options:nil error:&err]; + if (!lib) { printf("MSL compile failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id pso = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"iq2xxs_gemv"] error:&err]; + if (!pso) { printf("pipeline failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + + id bw = [dev newBufferWithBytes:w.data() length:w.size() options:MTLResourceStorageModeShared]; + id bx = [dev newBufferWithBytes:x.data() length:k * 4 options:MTLResourceStorageModeShared]; + id by = [dev newBufferWithLength:n * 4 options:MTLResourceStorageModeShared]; + id q = [dev newCommandQueue]; + unsigned uk = k, un = n; + double best = 1e9; + mach_timebase_info_data_t tb; mach_timebase_info(&tb); + for (int it = 0; it < 5; ++it) { + id cb = [q commandBuffer]; + id ce = [cb computeCommandEncoder]; + [ce setComputePipelineState:pso]; + [ce setBuffer:bw offset:0 atIndex:0]; + [ce setBuffer:bx offset:0 atIndex:1]; + [ce setBuffer:by offset:0 atIndex:2]; + [ce setBytes:&uk length:4 atIndex:3]; + [ce setBytes:&un length:4 atIndex:4]; + [ce dispatchThreads:MTLSizeMake((NSUInteger) n * 32, 1, 1) threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + [ce endEncoding]; + const uint64_t t0 = mach_absolute_time(); + [cb commit]; + [cb waitUntilCompleted]; + const double dt = (double) (mach_absolute_time() - t0) * tb.numer / tb.denom / 1e9; + if (dt < best) best = dt; + } + printf("throughput: %.2f GB/s (%.3f ms for %.1f MB)\n", (double) w.size() / best / 1e9, best * 1e3, + (double) w.size() / 1e6); + + // the CPU reference: the same decode, from the same tables + const float* gy = (const float*) [by contents]; + double num = 0, den = 0; + for (int r = 0; r < n; ++r) { + double acc = 0; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = &w[((size_t) r * nb + b) * bs]; + _Float16 dh; + memcpy(&dh, blk, 2); + const double d = (double) (float) dh; + for (int ib = 0; ib < 8; ++ib) { + uint32_t aux[2]; + memcpy(aux, blk + 2 + 4 * ib, 8); + const double db = d * (0.5 + (double) (aux[1] >> 28)) * 0.25; + const uint8_t* a8 = (const uint8_t*) aux; + for (int l = 0; l < 4; ++l) { + const uint64_t g = grid[a8[l]]; + const uint32_t sg = (uint32_t) signs_t[(aux[1] >> (7 * l)) & 127]; + const int xoff = ib * 32 + l * 8; + for (int j = 0; j < 8; ++j) { + const double wv = db * (double) ((g >> (8 * j)) & 0xff) * ((sg & kmask[j]) ? -1.0 : 1.0); + acc += wv * (double) x[xoff + j]; + } + } + } + } + const double diff = (double) gy[r] - acc; + num += diff * diff; + den += acc * acc; + } + const double l2 = sqrt(num / den); + printf("relative L2 %.3e\n", l2); + printf("%s\n", l2 < 1e-4 ? "PASS" : "FAIL"); + } + return 0; +} diff --git a/edge0-strata/metal/tests/dsv41/iq2xxs_t8_parity.mm b/edge0-strata/metal/tests/dsv41/iq2xxs_t8_parity.mm new file mode 100644 index 00000000..52085f24 --- /dev/null +++ b/edge0-strata/metal/tests/dsv41/iq2xxs_t8_parity.mm @@ -0,0 +1,186 @@ +// iq2xxs_t8_probe.mm - the tiled IQ2_XXS expert kernel: one weight decode feeding 8 tokens (the engine's 8-token window). +// +// The warp-per-row kernel measured 87.7 GB/s, which caps decode near 38 tok/s; the decode is the cost, so it must be paid once +// per window instead of once per token. Here each lane decodes its blocks ONCE and applies them to 8 activations, keeping eight +// accumulators. Success is a per-token effective rate well above the single-token one. +// clang++ -O2 -fobjc-arc -framework Metal -framework Foundation -o iq2xxs_t8_probe iq2xxs_t8_probe.mm && ./iq2xxs_t8_probe +#import +#import +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static std::vector load_msl_table(const std::string& path, const std::string& name, size_t expect) { + std::ifstream f(path); + std::string all((std::istreambuf_iterator(f)), std::istreambuf_iterator()); + std::regex re("constant \\w+ " + name + "\\[[^\\]]*\\] = \\{([^}]*)\\}"); + std::smatch m; + std::vector out; + if (!std::regex_search(all, m, re)) { fprintf(stderr, "table %s not found\n", name.c_str()); exit(2); } + const std::string body = m[1]; + size_t i = 0; + while (i < body.size() && out.size() < expect) { + while (i < body.size() && !isalnum((unsigned char) body[i])) ++i; + size_t j = i; + while (j < body.size() && (isalnum((unsigned char) body[j]) || body[j] == 'x')) ++j; + if (j > i) out.push_back(strtoull(body.substr(i, j - i).c_str(), nullptr, 0)); + i = j; + } + if (out.size() != expect) { fprintf(stderr, "table %s: %zu of %zu\n", name.c_str(), out.size(), expect); exit(2); } + return out; +} + +int main() { + @autoreleasepool { + const int n = 65536, k = 5120, T = 8; + const int nb = k / 256; + const size_t bs = 66; + auto grid = load_msl_table("ggml_grids.h", "iq2xxs_grid", 256); + auto signs_t = load_msl_table("ggml_grids.h", "ksigns_iq2xs", 128); + + std::vector w((size_t) n * nb * bs); + srand(3); + for (size_t i = 0; i < w.size(); ++i) w[i] = (uint8_t) (rand() & 0xff); + for (int r = 0; r < n; ++r) + for (int b = 0; b < nb; ++b) { + _Float16 d = (_Float16) (((rand() % 900) + 100) / 8000.0f); + memcpy(&w[((size_t) r * nb + b) * bs], &d, 2); + } + std::vector x((size_t) T * k), y((size_t) T * n); + for (size_t i = 0; i < x.size(); ++i) x[i] = ((rand() % 20001) - 10000) / 8192.0f; + + std::string grids; + { + std::ifstream f("ggml_grids.h"); + std::string line; + while (std::getline(f, line)) { + if (line.rfind("#pragma once", 0) == 0 || line.rfind("//", 0) == 0 || line.empty()) continue; + grids += line + "\n"; + } + } + std::string src = "#include \nusing namespace metal;\n" + grids + R"MSL( +constant uint kTokens = 8u; + +kernel void iq2xxs_t8(device const uchar* w [[buffer(0)]], device const float* x [[buffer(1)]], + device float* y [[buffer(2)]], constant uint& k [[buffer(3)]], constant uint& n [[buffer(4)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + const uint r = gid / 32u; + if (r >= n) return; + const uint nb = k / 256u; + const device uchar* base = w + (size_t) r * nb * 66u; + float acc[8]; + for (int t = 0; t < 8; ++t) acc[t] = 0.0f; + for (uint b = lane; b < nb; b += 32u) { + const device uchar* blk = base + (size_t) b * 66u; + const float d = float(as_type(*(const device ushort*) blk)); + for (int ib = 0; ib < 8; ++ib) { + const uint aux0 = uint(blk[2 + 4 * ib]) | (uint(blk[3 + 4 * ib]) << 8) | (uint(blk[4 + 4 * ib]) << 16) | (uint(blk[5 + 4 * ib]) << 24); + const uint aux1 = uint(blk[6 + 4 * ib]) | (uint(blk[7 + 4 * ib]) << 8) | (uint(blk[8 + 4 * ib]) << 16) | (uint(blk[9 + 4 * ib]) << 24); + const float db = d * (0.5f + float(aux1 >> 28)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const ulong g = iq2xxs_grid[(aux0 >> (8 * l)) & 0xff]; + const uint sg = uint(ksigns_iq2xs[(aux1 >> (7 * l)) & 127]); + const uint xoff = (uint) ib * 32u + (uint) l * 8u; + float wv[8]; + for (int j = 0; j < 8; ++j) { + wv[j] = db * float((g >> (8 * j)) & 0xff) * ((sg & kmask_iq2xs[j]) ? -1.0f : 1.0f); + } + for (uint t = 0; t < 8u; ++t) { + const device float* xt = x + (size_t) t * k + xoff; + float s = 0.0f; + for (int j = 0; j < 8; ++j) s += wv[j] * xt[j]; + acc[t] += s; + } + } + } + } + for (uint t = 0; t < 8u; ++t) { + float a = acc[t]; + a += simd_shuffle_xor(a, 16); + a += simd_shuffle_xor(a, 8); + a += simd_shuffle_xor(a, 4); + a += simd_shuffle_xor(a, 2); + a += simd_shuffle_xor(a, 1); + if (lane == 0) y[(size_t) t * n + r] = a; + } +} +)MSL"; + + id dev = MTLCreateSystemDefaultDevice(); + NSError* err = nil; + id lib = [dev newLibraryWithSource:[NSString stringWithUTF8String:src.c_str()] options:nil error:&err]; + if (!lib) { printf("MSL compile failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id pso = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"iq2xxs_t8"] error:&err]; + if (!pso) { printf("pipeline failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id bw = [dev newBufferWithBytes:w.data() length:w.size() options:MTLResourceStorageModeShared]; + id bx = [dev newBufferWithBytes:x.data() length:x.size() * 4 options:MTLResourceStorageModeShared]; + id by = [dev newBufferWithLength:y.size() * 4 options:MTLResourceStorageModeShared]; + id q = [dev newCommandQueue]; + unsigned uk = k, un = n; + double best = 1e9; + mach_timebase_info_data_t tb; mach_timebase_info(&tb); + for (int it = 0; it < 5; ++it) { + id cb = [q commandBuffer]; + id ce = [cb computeCommandEncoder]; + [ce setComputePipelineState:pso]; + [ce setBuffer:bw offset:0 atIndex:0]; + [ce setBuffer:bx offset:0 atIndex:1]; + [ce setBuffer:by offset:0 atIndex:2]; + [ce setBytes:&uk length:4 atIndex:3]; + [ce setBytes:&un length:4 atIndex:4]; + [ce dispatchThreads:MTLSizeMake((NSUInteger) n * 32, 1, 1) threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + [ce endEncoding]; + const uint64_t t0 = mach_absolute_time(); + [cb commit]; + [cb waitUntilCompleted]; + const double dt = (double) (mach_absolute_time() - t0) * tb.numer / tb.denom / 1e9; + if (dt < best) best = dt; + } + // the window reads the weights once and produces 8 tokens of output + printf("window of 8: %.2f GB/s effective (%.2f GB/s raw), %.3f ms for %.1f MB\n", 8.0 * w.size() / best / 1e9, + (double) w.size() / best / 1e9, best * 1e3, (double) w.size() / 1e6); + + const float* gy = (const float*) [by contents]; + const uint8_t* kmask = (const uint8_t*) "\x01\x02\x04\x08\x10\x20\x40\x80"; + double num = 0, den = 0; + for (int r = 0; r < n; r += 97) { // sampled: the reference is O(T * n * k) + for (int t = 0; t < T; ++t) { + double acc = 0; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = &w[((size_t) r * nb + b) * bs]; + _Float16 dh; memcpy(&dh, blk, 2); + const double d = (double) (float) dh; + for (int ib = 0; ib < 8; ++ib) { + uint32_t aux[2]; + memcpy(aux, blk + 2 + 4 * ib, 8); + const double db = d * (0.5 + (double) (aux[1] >> 28)) * 0.25; + const uint8_t* a8 = (const uint8_t*) aux; + for (int l = 0; l < 4; ++l) { + const uint64_t g = grid[a8[l]]; + const uint32_t sg = (uint32_t) signs_t[(aux[1] >> (7 * l)) & 127]; + const int xoff = ib * 32 + l * 8; + for (int j = 0; j < 8; ++j) { + const double wv = db * (double) ((g >> (8 * j)) & 0xff) * ((sg & kmask[j]) ? -1.0 : 1.0); + acc += wv * (double) x[(size_t) t * k + xoff + j]; + } + } + } + } + const double diff = (double) gy[(size_t) t * n + r] - acc; + num += diff * diff; + den += acc * acc; + } + } + const double l2 = sqrt(num / den); + printf("relative L2 %.3e (sampled rows)\n", l2); + printf("%s\n", l2 < 1e-5 ? "PASS" : "FAIL"); + } + return 0; +} diff --git a/edge0-strata/metal/tests/dsv41/q8_gemv_parity.mm b/edge0-strata/metal/tests/dsv41/q8_gemv_parity.mm new file mode 100644 index 00000000..6eda27bd --- /dev/null +++ b/edge0-strata/metal/tests/dsv41/q8_gemv_parity.mm @@ -0,0 +1,117 @@ +// q8gemv_probe.mm - the quantised-weight GEMV in MSL: Q8_0 weights times F32 activations (ds4_gguf.cu's Q8 path). +// +// Q8_0 block: 32 int8 values + one f16 scale d, 34 bytes. y[r] = sum over blocks of d * sum_i q[i] * x[i]. This is the +// simplest quantised-weight kernel and the one every Q8 tensor in the pack uses; IQ2/IQ3 use the same shape with a grid +// table instead of a plain int8 array, so this validates the code path before the harder formats. +// clang++ -O2 -fobjc-arc -framework Metal -framework Foundation -o q8gemv_probe q8gemv_probe.mm && ./q8gemv_probe +#import +#import +#include +#include +#include +#include +#include +#include + +static const char* kSrc = R"MSL( +#include +using namespace metal; + +// one warp a row, lane-strided over the blocks +kernel void ds4_q8_gemv(device const char* q [[buffer(0)]], device const half* d [[buffer(1)]], + device const float* x [[buffer(2)]], device float* y [[buffer(3)]], + constant uint& k [[buffer(4)]], constant uint& n [[buffer(5)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + const uint r = gid / 32u; + if (r >= n) return; + const uint nb = k / 32u; // blocks a row + const device char* qr = q + (size_t) r * k; // k int8 values, one a float + const device half* dr = d + (size_t) r * nb; + float acc = 0.0f; + for (uint b = lane; b < nb; b += 32u) { + float s = 0.0f; + for (uint i = 0; i < 32u; ++i) s += float(qr[b * 32u + i]) * x[b * 32u + i]; + acc += s * float(dr[b]); + } + acc += simd_shuffle_xor(acc, 16); + acc += simd_shuffle_xor(acc, 8); + acc += simd_shuffle_xor(acc, 4); + acc += simd_shuffle_xor(acc, 2); + acc += simd_shuffle_xor(acc, 1); + if (lane == 0) y[r] = acc; +} +)MSL"; + +int main() { + @autoreleasepool { + const int n = 16384, k = 5120; + const int nb = k / 32; + std::vector q((size_t) n * k); + std::vector dv((size_t) n * nb), x(k), y(n); + srand(5); + for (size_t i = 0; i < q.size(); ++i) q[i] = (signed char) ((rand() % 255) - 127); + for (size_t i = 0; i < dv.size(); ++i) dv[i] = ((rand() % 1000) + 1) / 4000.0f; + for (int i = 0; i < k; ++i) x[i] = ((rand() % 20001) - 10000) / 8192.0f; + + // the f16 scales as half bits + std::vector dh(dv.size()); + for (size_t i = 0; i < dv.size(); ++i) { + _Float16 h = (_Float16) dv[i]; + memcpy(&dh[i], &h, 2); + } + + id dev = MTLCreateSystemDefaultDevice(); + NSError* err = nil; + id lib = [dev newLibraryWithSource:[NSString stringWithUTF8String:kSrc] options:nil error:&err]; + if (!lib) { printf("MSL compile failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id pso = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"ds4_q8_gemv"] error:&err]; + if (!pso) { printf("pipeline failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id bq = [dev newBufferWithBytes:q.data() length:q.size() options:MTLResourceStorageModeShared]; + id bd = [dev newBufferWithBytes:dh.data() length:dh.size() * 2 options:MTLResourceStorageModeShared]; + id bx = [dev newBufferWithBytes:x.data() length:k * 4 options:MTLResourceStorageModeShared]; + id by = [dev newBufferWithLength:n * 4 options:MTLResourceStorageModeShared]; + id qq = [dev newCommandQueue]; + unsigned uk = k, un = n; + double best = 1e9; + mach_timebase_info_data_t tb; mach_timebase_info(&tb); + for (int it = 0; it < 6; ++it) { + id cb = [qq commandBuffer]; + id ce = [cb computeCommandEncoder]; + [ce setComputePipelineState:pso]; + [ce setBuffer:bq offset:0 atIndex:0]; + [ce setBuffer:bd offset:0 atIndex:1]; + [ce setBuffer:bx offset:0 atIndex:2]; + [ce setBuffer:by offset:0 atIndex:3]; + [ce setBytes:&uk length:4 atIndex:4]; + [ce setBytes:&un length:4 atIndex:5]; + [ce dispatchThreads:MTLSizeMake((NSUInteger) n * 32, 1, 1) threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + [ce endEncoding]; + const uint64_t t0 = mach_absolute_time(); + [cb commit]; + [cb waitUntilCompleted]; + const double dt = (double) (mach_absolute_time() - t0) * tb.numer / tb.denom / 1e9; + if (dt < best) best = dt; + } + const double bytes = (double) q.size() + (double) dh.size() * 2.0; + printf("throughput: %.2f GB/s, %.2f GFLOP/s (%.3f ms for %.1f MB)\n", bytes / best / 1e9, 2.0 * n * k / best / 1e9, + best * 1e3, bytes / 1e6); + + const float* gy = (const float*) [by contents]; + double num = 0, den = 0; + for (int r = 0; r < n; ++r) { + double acc = 0; + for (int b = 0; b < nb; ++b) { + double s = 0; + for (int i = 0; i < 32; ++i) s += (double) q[(size_t) r * k + b * 32 + i] * (double) x[b * 32 + i]; + acc += s * (double) (float) (_Float16) dv[(size_t) r * nb + b]; // the kernel gets the half-rounded scale + } + const double diff = (double) gy[r] - acc; + num += diff * diff; + den += acc * acc; + } + const double l2 = sqrt(num / den); + printf("relative L2 %.3e\n", l2); + printf("%s\n", l2 < 1e-5 ? "PASS" : "FAIL"); + } + return 0; +} diff --git a/edge0-strata/metal/tests/dsv41/q8_quant_parity.mm b/edge0-strata/metal/tests/dsv41/q8_quant_parity.mm new file mode 100644 index 00000000..d46a86ff --- /dev/null +++ b/edge0-strata/metal/tests/dsv41/q8_quant_parity.mm @@ -0,0 +1,105 @@ +// q8probe.mm - the first dsv41 kernel in MSL: the activation quantiser at ds4_gguf.cu:201. +// +// CUDA's q8 path: per 32 values, d = max|x| / 127 (a half), then q = round(x / d) clamped to [-128, 127], stored as int8 +// with the f16 scale. This compiles the MSL version at run time, runs it on the M5 Max GPU, and compares it against the +// same arithmetic in C++ - if the two agree, this function is ready to become a Pipeline in metal/src/engine. +// clang++ -O2 -fobjc-arc -framework Metal -framework Foundation -o q8probe q8probe.mm && ./q8probe +#import +#import +#include +#include +#include +#include +#include + +static const char* kSrc = R"MSL( +#include +using namespace metal; + +// one threadgroup of 32 threads per block of 32 values (CUDA's layout), one block per 32 inputs +kernel void q8_quant(device const float* x [[buffer(0)]], constant uint& ldx [[buffer(1)]], + device char* q [[buffer(2)]], device half* d [[buffer(3)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_threadgroup]]) { + const uint b = gid / 32u; // the block this thread belongs to + const int base = int(b) * int(ldx); + float v = x[base + int(lane)]; + float amax = fabs(v); + amax = max(amax, simd_shuffle_xor(amax, 16)); + amax = max(amax, simd_shuffle_xor(amax, 8)); + amax = max(amax, simd_shuffle_xor(amax, 4)); + amax = max(amax, simd_shuffle_xor(amax, 2)); + amax = max(amax, simd_shuffle_xor(amax, 1)); + const float scale = amax / 127.0f; + const float inv = scale > 0.0f ? 1.0f / scale : 0.0f; + int iq = int(rint(v * inv)); + iq = clamp(iq, -128, 127); + q[int(b) * 32 + int(lane)] = char(iq); + if (lane == 0) d[int(b)] = half(scale); +} +)MSL"; + +int main() { + @autoreleasepool { + const int n = 32 * 4096; + std::vector x(n); + srand(7); + for (int i = 0; i < n; ++i) x[i] = ((rand() % 20001) - 10000) / 997.0f; + + id dev = MTLCreateSystemDefaultDevice(); + NSError* err = nil; + id lib = [dev newLibraryWithSource:[NSString stringWithUTF8String:kSrc] options:nil error:&err]; + if (!lib) { printf("MSL compile failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + id pso = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"q8_quant"] error:&err]; + if (!pso) { printf("pipeline failed: %s\n", [[err localizedDescription] UTF8String]); return 1; } + + id bx = [dev newBufferWithBytes:x.data() length:n * sizeof(float) options:MTLResourceStorageModeShared]; + id bq = [dev newBufferWithLength:n options:MTLResourceStorageModeShared]; + id bd = [dev newBufferWithLength:(n / 32) * 2 options:MTLResourceStorageModeShared]; + id qq = [dev newCommandQueue]; + id cb = [qq commandBuffer]; + id ce = [cb computeCommandEncoder]; + [ce setComputePipelineState:pso]; + [ce setBuffer:bx offset:0 atIndex:0]; + unsigned ldx = 32; // the stride between blocks of 32 values + [ce setBytes:&ldx length:4 atIndex:1]; + [ce setBuffer:bq offset:0 atIndex:1 + 1]; + [ce setBuffer:bd offset:0 atIndex:2 + 1]; + [ce dispatchThreads:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + [ce endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + + // the same arithmetic on the CPU + const signed char* gq = (const signed char*) [bq contents]; + const _Float16* gd = (const _Float16*) [bd contents]; + // first block, side by side + { + float amax = 0; for (int i = 0; i < 32; ++i) amax = fmaxf(amax, fabsf(x[i])); + const float scale = amax / 127.0f, inv = scale > 0 ? 1.0f / scale : 0.0f; + printf("block0: gpu d=%.6f cpu d=%.6f\n", (double) gd[0], (double) scale); + printf(" gpu q:"); for (int i = 0; i < 8; ++i) printf(" %d", (int) gq[i]); + printf("\n cpu q:"); for (int i = 0; i < 8; ++i) printf(" %d", (int) rintf(x[i] * inv)); + printf("\nblock1: gpu d=%.6f gpu q:", (double) gd[1]); for (int i = 32; i < 40; ++i) printf(" %d", (int) gq[i]); + float am1 = 0; for (int i = 32; i < 64; ++i) am1 = fmaxf(am1, fabsf(x[i])); + printf("\n cpu d=%.6f cpu q:", (double) (am1 / 127.0f)); for (int i = 32; i < 40; ++i) printf(" %d", (int) rintf(x[i] * (127.0f / am1))); + printf("\n"); + } int bad = 0; double worst = 0; + for (int b = 0; b < n / 32; ++b) { + float amax = 0; + for (int i = 0; i < 32; ++i) amax = fmaxf(amax, fabsf(x[b * 32 + i])); + const float scale = amax / 127.0f; + const float inv = scale > 0 ? 1.0f / scale : 0.0f; + const double dd = fabs((double) gd[b] - (double) scale) / (scale > 0 ? scale : 1); + worst = fmax(worst, dd); + for (int i = 0; i < 32; ++i) { + int iq = (int) rintf(x[b * 32 + i] * inv); + iq = iq < -128 ? -128 : (iq > 127 ? 127 : iq); + if (iq != gq[b * 32 + i]) ++bad; + } + } + printf("blocks %d, values %d\n", n / 32, n); + printf("int8 differ: %d worst scale rel err %.3e\n", bad, worst); + printf("%s\n", bad == 0 && worst < 1e-3 ? "PASS" : "FAIL"); + } + return 0; +} diff --git a/edge0-strata/metal/tests/expert_parity.cpp b/edge0-strata/metal/tests/expert_parity.cpp new file mode 100644 index 00000000..5a596a74 --- /dev/null +++ b/edge0-strata/metal/tests/expert_parity.cpp @@ -0,0 +1,669 @@ +// metal/tests/expert_parity.cpp - the grouped routed-expert kernels (moe.metal) against a double-precision CPU +// reference, and their speed. +// +// expert_parity [--gu Q2_0] [--dn Q2_0] [--ntok 1..8] [--gucfg L,R,S] [--dncfg L,R,S] parity (synthetic) +// expert_parity --bench [--gu ..] [--dn ..] [--ntok N] [--layers 8] [--reps 96] [--sweep] speed +// +// Gyro (agentionai's trellis experts, TQ2_T / TQK6 / TQK7): +// expert_parity --gu TQK6 --dn TQK7 [--had-pack ] synthetic blocks; with a pack, the Hadamard rotation of +// x (CPU) and of h (the Metal had128 kernel) is part of it +// expert_parity --gguf --layer L [--had-pack ] [--ntok N] [--seed S] REAL experts of layer L (64 of +// the 512, random) instead of random blocks +// expert_parity --had-test [--had-pack ] the Metal had128 kernels (f32, f16, f16 rows) against the CPU rotation, bitwise +// +// Parity: random valid blocks for 64 experts, random activations x ~ U(-1,1), random distinct top-10 per token, +// weights summing to 1. Checked: the plan (groups by expert id, entries in (token, k) order), h = silu(gate) * up +// per entry and y = sum_k w * down . h, each against double; reported as max |err| / max |ref| and relative L2. +// Bench: `layers` distinct copies of a 512-expert arena (no reuse across reps beyond chance; the SLC is ~32 MB), +// a different random routing per rep; GB/s = bytes of the distinct experts the kernel must read / GPU time. +#include "mtl_rt.h" +#include "ref_quant.h" +#include "gguf_lite.h" +// Gyro: the engine's CPU Hadamard rotation (the reference of the Metal one) and hadamard.txt's parser +#include "../../src/kernels/cpu/hadamard.cpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ggml-common-tq.h" +const uint16_t* tq_lut_for_test() { return tq_lut_f16; } + +namespace { + +constexpr int N_EMBD = 2560, N_FF = 640, TOPK = 10; +// plan layout (moe.metal) +constexpr int P_NG = 0, P_IGU = 4, P_IDN = 8, P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, + P_ETOK = P_GSTART + P_MAXE + 4, P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, + P_INTS = P_TOKENT + P_MAXE; + +struct MoeArgs { + uint64_t gate_off, up_off, down_off, gu_stride, d_stride; + uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad; +}; + +int type_of(const std::string& s) { + for (int t : {ref::TQ2_T, ref::TQK6, ref::TQK7, ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::Q8_0, ref::Q5_1}) + if (s == ref::type_name(t)) return t; + std::fprintf(stderr, "unknown type %s\n", s.c_str()); + std::exit(2); +} + +struct Cfg { int L, R, S; }; +Cfg parse_cfg(const char* s) { + Cfg c{}; + if (std::sscanf(s, "%d,%d,%d", &c.L, &c.R, &c.S) != 3) { std::fprintf(stderr, "bad cfg %s\n", s); std::exit(2); } + return c; +} +std::string kname(const char* k, int t, Cfg c) { + return std::string(k) + "_" + std::to_string(t) + "_" + std::to_string(c.L) + "_" + std::to_string(c.R) + "_" + + std::to_string(c.S); +} +int rows_per_tg(Cfg c) { return c.S * (32 / c.L) * c.R; } + +struct Layout { + int gu, dn; + size_t gu_row, d_row, up_off, down_off, blob; +}; +Layout layout(int gu, int dn) { + Layout L{gu, dn}; + L.gu_row = ref::row_bytes(gu, N_EMBD); + L.d_row = ref::row_bytes(dn, N_FF); + L.up_off = (size_t)N_FF * L.gu_row; + L.down_off = 2 * L.up_off; + L.blob = L.down_off + (size_t)N_EMBD * L.d_row; + return L; +} + +void random_routing(ref::Rng& r, int ntok, int n_expert, std::vector& ids, std::vector& w, + double share = 0.0, const std::vector* prev = nullptr) { + ids.assign(ntok * TOPK, 0); + w.assign(ntok * TOPK, 0.0f); + for (int t = 0; t < ntok; ++t) { + std::set used; + float sum = 0; + for (int k = 0; k < TOPK; ++k) { + int e; + do { + // with probability `share`, reuse an expert of the previous token (verify windows share experts) + if (t > 0 && share > 0 && r.uni() < share) e = ids[(t - 1) * TOPK + (r.next() % TOPK)]; + else e = (int)(r.next() % n_expert); + } while (used.count(e)); + used.insert(e); + ids[t * TOPK + k] = e; + w[t * TOPK + k] = 0.2f + r.uni(); + sum += w[t * TOPK + k]; + } + for (int k = 0; k < TOPK; ++k) w[t * TOPK + k] /= sum; + } + (void)prev; +} + +struct Kernels { + mtl::Pipeline plan, gu, dn, comb, prep; + Cfg cgu, cdn; +}; +int xkind(int t) { + switch (t) { + case ref::Q2_0: return 1; + case ref::Q4_K: case ref::Q5_K: return 2; + case ref::Q6_K: case ref::Q3_K: return 3; + default: return 0; + } +} +constexpr int CAPE = 8 * TOPK; // entries capacity of h +// prepared buffers: x [8][N_EMBD] -> XPX, h [CAPE][N_FF] -> XPH (4 planes + S16) +constexpr size_t PLANE_X = 8 * N_EMBD, PLANE_H = (size_t)CAPE * N_FF; + +mtl::Dispatch prep_dispatch(const mtl::Pipeline& p, const mtl::Buffer& x, const mtl::Buffer& xp, uint32_t n, uint32_t nt) { + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(x), mtl::Arg::B(xp), mtl::Arg::V(n), mtl::Arg::V(nt)}; + d.tg[0] = 256; + d.grid[0] = (n / 16 * nt + 255) / 256; + return d; +} + +int g_gu_type = 0, g_dn_type = 0; +// Gyro: the rotation of h between the gate/up and the down dispatch (the Metal had128_f32 kernel, signs of width n_ff) +struct HadArgs { uint32_t width, x_stride, y_stride; float scale; uint32_t has_signs, pad; }; +const mtl::Pipeline* g_had32 = nullptr; +const mtl::Buffer* g_hsh = nullptr; +bool g_had_h = false; +bool g_dn_prep = false; // Q2_0 down: prepared h (an xprep dispatch) instead of the raw-x decoder +const mtl::Buffer* g_xpx = nullptr; +const mtl::Buffer* g_xph = nullptr; + +void encode_layer(std::vector& ds, const Kernels& K, const mtl::Buffer& W, size_t w_off, + const mtl::Buffer& x, const mtl::Buffer& h, const mtl::Buffer& ye, const mtl::Buffer& y, + const mtl::Buffer& ids, const mtl::Buffer& wts, const mtl::Buffer& plan, size_t plan_off, + const MoeArgs& a, bool with_plan, bool with_gu, bool with_dn, bool with_comb) { + const mtl::Buffer& XPX = *g_xpx; + const mtl::Buffer& XPH = *g_xph; + if (with_plan) { + mtl::Dispatch d; + d.p = &K.plan; + d.args = {mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan, plan_off), mtl::Arg::V(a)}; + d.tg[0] = 512; + ds.push_back(d); + } + if (with_gu) { + mtl::Dispatch d; + d.p = &K.gu; + d.args = {mtl::Arg::B(W, w_off), mtl::Arg::B(XPX, 4 * PLANE_X * xkind(g_gu_type)), mtl::Arg::B(h), + mtl::Arg::B(plan, plan_off), mtl::Arg::V(a), mtl::Arg::B(XPX, 4 * PLANE_X * 4)}; + d.tg[0] = K.cgu.S * 32; + d.indirect = &plan; + d.ind_off = plan_off + P_IGU * 4; + ds.push_back(d); + if (g_had_h) { // H h in place, one row per entry (the entries are the window's token x k) + static HadArgs ha; + ha = HadArgs{(uint32_t)N_FF, (uint32_t)N_FF, (uint32_t)N_FF, 1.0f / std::sqrt(128.0f), g_hsh ? 1u : 0u, 0u}; + mtl::Dispatch hd; + hd.p = g_had32; + hd.args = {mtl::Arg::B(h), mtl::Arg::B(h), mtl::Arg::B(g_hsh ? *g_hsh : h), mtl::Arg::V(ha)}; + hd.grid[0] = N_FF / 128; + hd.grid[1] = a.n_tok * TOPK; + hd.tg[0] = 128; + ds.push_back(hd); + } + } + if (with_dn) { + const int dk = g_dn_prep ? xkind(g_dn_type) : 0; + if (dk) ds.push_back(prep_dispatch(K.prep, h, XPH, N_FF, CAPE)); + mtl::Dispatch d; + d.p = &K.dn; + d.args = {mtl::Arg::B(W, w_off), + dk ? mtl::Arg::B(XPH, 4 * PLANE_H * dk) : mtl::Arg::B(h), mtl::Arg::B(ye), + mtl::Arg::B(plan, plan_off), mtl::Arg::V(a), mtl::Arg::B(XPH, 4 * PLANE_H * 4)}; + d.tg[0] = K.cdn.S * 32; + d.indirect = &plan; + d.ind_off = plan_off + P_IDN * 4; + ds.push_back(d); + } + if (with_comb) { + mtl::Dispatch d; + d.p = &K.comb; + d.args = {mtl::Arg::B(ye), mtl::Arg::B(y), mtl::Arg::B(plan, plan_off), mtl::Arg::V(a)}; + d.tg[0] = 256; + d.grid[0] = (a.n_tok * N_EMBD + 255) / 256; + ds.push_back(d); + } +} + +struct Err { + double max_abs = 0, max_ref = 0, se = 0, sr = 0; + void add(double got, double ref) { + max_abs = std::max(max_abs, std::fabs(got - ref)); + max_ref = std::max(max_ref, std::fabs(ref)); + se += (got - ref) * (got - ref); + sr += ref * ref; + } + double rel_max() const { return max_ref > 0 ? max_abs / max_ref : max_abs; } + double rel_l2() const { return sr > 0 ? std::sqrt(se / sr) : std::sqrt(se); } +}; + +double median(std::vector v) { + std::sort(v.begin(), v.end()); + return v[v.size() / 2]; +} + +} // namespace + +int main(int argc, char** argv) { + int gu = ref::Q2_0, dn = ref::Q2_0, ntok = 1, layers = 8, reps = 96; + Cfg cgu{32, 2, 4}, cdn{8, 2, 4}; + bool bench = false, sweep = false, had_test = false, decode_check = false; + double share = 0.0; + std::string gguf_path, had_pack; + int real_layer = 0, experts_first = 0; + uint64_t seed = 1; + bool cfg_gu_set = false, cfg_dn_set = false; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + auto nx = [&] { return i + 1 < argc ? argv[++i] : (std::fprintf(stderr, "%s needs a value\n", a.c_str()), std::exit(2), (char*)nullptr); }; + if (a == "--gu") gu = type_of(nx()); + else if (a == "--dn") dn = type_of(nx()); + else if (a == "--ntok") ntok = std::atoi(nx()); + else if (a == "--gucfg") { cgu = parse_cfg(nx()); cfg_gu_set = true; } + else if (a == "--dncfg") { cdn = parse_cfg(nx()); cfg_dn_set = true; } + else if (a == "--bench") bench = true; + else if (a == "--sweep") sweep = true; + else if (a == "--layers") layers = std::atoi(nx()); + else if (a == "--reps") reps = std::atoi(nx()); + else if (a == "--share") share = std::atof(nx()); + else if (a == "--dn-prep") g_dn_prep = true; + else if (a == "--gguf") gguf_path = nx(); + else if (a == "--layer") real_layer = std::atoi(nx()); + else if (a == "--experts-first") experts_first = std::atoi(nx()); // only experts 0..N-1 of the layer are in the file (a sparse copy) + else if (a == "--had-pack") had_pack = nx(); + else if (a == "--had-test") had_test = true; + else if (a == "--decode-check") decode_check = true; + else if (a == "--seed") seed = std::strtoull(nx(), nullptr, 10); + else { std::fprintf(stderr, "unknown argument %s\n", a.c_str()); return 2; } + } + (void)cfg_gu_set; + (void)cfg_dn_set; + if (dn == ref::IQ4_NL && cdn.L == 8 && !cfg_dn_set) cdn = {4, 2, 4}; + gguf::File gf; + const gguf::Tensor *tg_ = nullptr, *tu_ = nullptr, *td_ = nullptr; + if (!gguf_path.empty()) { // real experts: the types come from the file + std::string err; + if (!gf.open(gguf_path, err)) { std::fprintf(stderr, "%s\n", err.c_str()); return 2; } + const std::string b = "blk." + std::to_string(real_layer) + "."; + tg_ = gf.find(b + "ffn_gate_exps.weight"); + tu_ = gf.find(b + "ffn_up_exps.weight"); + td_ = gf.find(b + "ffn_down_exps.weight"); + if (!tg_ || !tu_ || !td_) { std::fprintf(stderr, "no routed experts in layer %d of %s\n", real_layer, gguf_path.c_str()); return 2; } + gu = tg_->type; + dn = td_->type; + } + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal"}); + const Layout L = layout(gu, dn); + // Gyro: the pack's hadamard.txt (the signs of the widths 2560 and 640) -> the engine's CPU spec (the reference) + mtl::Buffer hsx, hsh; + mtl::Pipeline had32 = cx.pipeline(lib, "had128_f32"), had16 = cx.pipeline(lib, "had128_f16"), had16r = cx.pipeline(lib, "had128_f16_rows"); + const float* sgx = nullptr; + const float* sgh = nullptr; + if (!had_pack.empty()) { + strata::kernels::cpu::HadamardSpec spec; + std::string err; + if (!strata::kernels::cpu::hadamard_parse(had_pack + "/hadamard.txt", 48, spec, err)) { + std::fprintf(stderr, "%s\n", err.c_str()); + return 2; + } + strata::kernels::cpu::hadamard_set(std::move(spec)); + sgx = strata::kernels::cpu::hadamard_signs(N_EMBD); + sgh = strata::kernels::cpu::hadamard_signs(N_FF); + if (sgx) { hsx = cx.alloc(4 * N_EMBD); std::memcpy(hsx.ptr, sgx, 4 * N_EMBD); } + if (sgh) { hsh = cx.alloc(4 * N_FF); std::memcpy(hsh.ptr, sgh, 4 * N_FF); } + g_had32 = &had32; + g_hsh = sgh ? &hsh : nullptr; + g_had_h = !had_test; + std::printf("Hadamard: block %d, signs for %s%s\n", strata::kernels::cpu::hadamard_spec().block, sgx ? "2560 " : "", sgh ? "640" : "(none)"); + } + if (decode_check) { + // the decoder the Metal kernels and ref_quant share the meaning of, written a third time straight from the format's + // definition (ggml-common-tq.h): step t's state = 16 stream bits from bit (31 - t) * K, LSB first, circular modulo 32 K + // bits, masked to 15 bits (TQ2_T: qs[(t + 31) % 32] << 8 | qs[t]); x = state * 0x9e3779b1; points lut[x >> 21] and + // lut[(x >> 10) & 2047]; weights = fp16 values * d. Compared bit for bit with ref::dequant on the real tensors. + extern const uint16_t* tq_lut_for_test(); + long blocks = 0, bad = 0; + const gguf::Tensor* ts[3] = {tg_, tu_, td_}; + if (!tg_) { std::fprintf(stderr, "--decode-check needs --gguf\n"); return 2; } + for (const gguf::Tensor* t : ts) { + const int ty = t->type, K = ty == ref::TQ2_T ? 8 : ty == ref::TQK6 ? 6 : 7, bb = ref::block_bytes(ty); + const uint8_t* base = (const uint8_t*)gf.map + t->offset; + const size_t nb = t->bytes / bb; + const size_t step = std::max(1, nb / 200000); // up to 200k blocks a tensor, spread over it + for (size_t i = 0; i < nb; i += step) { + const uint8_t* b = base + i * bb; + float want[128], got[128]; + const float d = ref::f16_to_f32((uint16_t)(b[0] | (b[1] << 8))); + const uint8_t* qs = b + 2; + for (int st = 0; st < 32; ++st) { + uint32_t state; + if (K == 8) state = (((uint32_t)qs[(st + 31) % 32] << 8) | qs[st]) & 0x7fff; + else { + const int off = ((31 - st) * K) % (32 * K); + uint32_t v = 0; + for (int bit = 0; bit < 16; ++bit) { + const int pos = (off + bit) % (32 * K); + v |= (uint32_t)((qs[pos / 8] >> (pos % 8)) & 1) << bit; + } + state = v & 0x7fff; + } + const uint32_t x = state * 0x9e3779b1u; + const uint16_t* lut = tq_lut_for_test(); + const uint32_t pa = x >> 21, pb = (x >> 10) & 2047u; + want[4 * st] = ref::f16_to_f32(lut[2 * pa]) * d; + want[4 * st + 1] = ref::f16_to_f32(lut[2 * pa + 1]) * d; + want[4 * st + 2] = ref::f16_to_f32(lut[2 * pb]) * d; + want[4 * st + 3] = ref::f16_to_f32(lut[2 * pb + 1]) * d; + } + ref::dequant(ty, b, got, 128); + bad += std::memcmp(want, got, sizeof want) != 0; + ++blocks; + } + } + std::printf("decode-check layer %d: %ld blocks of the real tensors, %ld differ from the format's definition %s\n", real_layer, + blocks, bad, bad == 0 ? "PASS" : "FAIL"); + return bad == 0 ? 0 : 1; + } + if (had_test) { + // the Metal rotation kernels against the CPU's (hadamard_rotate), bit for bit + int bad = 0; + ref::Rng r2(77 + seed); + for (int width : {N_FF, N_EMBD}) { + const int rows = 37; + const float* sg = width == N_EMBD ? sgx : sgh; + const mtl::Buffer& sb = width == N_EMBD ? hsx : hsh; + const float scale = 1.0f / std::sqrt(128.0f); + std::vector in((size_t)rows * width); + for (auto& v : in) v = r2.sym() * 3.0f; + mtl::Buffer bx = cx.alloc(4 * in.size()), by = cx.alloc(4 * in.size()); + std::memcpy(bx.ptr, in.data(), 4 * in.size()); + HadArgs ha{(uint32_t)width, (uint32_t)width, (uint32_t)width, scale, sg ? 1u : 0u, 0u}; + auto run32 = [&](const mtl::Buffer& src, const mtl::Buffer& dst) { + mtl::Dispatch d; + d.p = &had32; + d.args = {mtl::Arg::B(src), mtl::Arg::B(dst), mtl::Arg::B(sg ? sb : src), mtl::Arg::V(ha)}; + d.grid[0] = width / 128; d.grid[1] = rows; d.tg[0] = 128; + cx.run({d}); + }; + run32(bx, by); // out of place + std::vector want = in; + for (int r = 0; r < rows; ++r) strata::kernels::cpu::hadamard_rotate(want.data() + (size_t)r * width, width, 128, sg); + int b32 = 0; + for (size_t i = 0; i < in.size(); ++i) b32 += std::memcmp(&want[i], by.as() + i, 4) != 0; + run32(bx, bx); // in place + int b32i = 0; + for (size_t i = 0; i < in.size(); ++i) b32i += std::memcmp(&want[i], bx.as() + i, 4) != 0; + // f16 (no signs): out of place, and the rows kernel in place; the CPU's float result rounded to half + std::vector<_Float16> hin(in.size()); + for (size_t i = 0; i < in.size(); ++i) hin[i] = (_Float16)in[i]; + std::vector hw(in.size()); + for (size_t i = 0; i < in.size(); ++i) hw[i] = (float)hin[i]; + for (int r = 0; r < rows; ++r) strata::kernels::cpu::hadamard_rotate(hw.data() + (size_t)r * width, width, 128, nullptr); + mtl::Buffer hx = cx.alloc(2 * in.size()), hy = cx.alloc(2 * in.size()); + std::memcpy(hx.ptr, hin.data(), 2 * in.size()); + HadArgs hb{(uint32_t)width, (uint32_t)width, (uint32_t)width, scale, 0u, 0u}; + { + mtl::Dispatch d; + d.p = &had16; + d.args = {mtl::Arg::B(hx), mtl::Arg::B(hy), mtl::Arg::B(hx), mtl::Arg::V(hb)}; + d.grid[0] = width / 128; d.grid[1] = rows; d.tg[0] = 128; + cx.run({d}); + } + int b16 = 0; + for (size_t i = 0; i < in.size(); ++i) { + const _Float16 w = (_Float16)hw[i]; + b16 += std::memcmp(&w, hy.as<_Float16>() + i, 2) != 0; + } + { + mtl::Dispatch d; + d.p = &had16r; + d.args = {mtl::Arg::B(hx), mtl::Arg::B(hx), mtl::Arg::B(hx), mtl::Arg::V(hb)}; + d.grid[0] = rows; d.tg[0] = 128; + cx.run({d}); + } + int b16r = 0; + for (size_t i = 0; i < in.size(); ++i) { + const _Float16 w = (_Float16)hw[i]; + b16r += std::memcmp(&w, hx.as<_Float16>() + i, 2) != 0; + } + std::printf("had-test width %d, %d rows%s: f32 %d / f32 in place %d / f16 %d / f16 rows %d of %zu values differ from the CPU\n", + width, rows, sg ? " (with signs)" : " (no signs)", b32, b32i, b16, b16r, in.size()); + bad += b32 + b32i + b16 + b16r; + } + std::printf("had-test %s\n", bad == 0 ? "PASS (bitwise)" : "FAIL"); + return bad == 0 ? 0 : 1; + } + std::printf("device %s; gate/up %s, down %s, blob %zu bytes; %d token(s)\n", cx.device_name().c_str(), + ref::type_name(gu), ref::type_name(dn), L.blob, ntok); + + auto make = [&](Cfg a, Cfg b) { + Kernels K; + K.plan = cx.pipeline(lib, "moe_plan"); + K.prep = cx.pipeline(lib, "xprep"); + K.comb = cx.pipeline(lib, "moe_combine"); + K.gu = cx.pipeline(lib, kname("moe_gu", gu, a), mtl::Constants().i(22, std::getenv("MOE_PROBE") ? std::atoi(std::getenv("MOE_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)); + K.dn = cx.pipeline(lib, kname("moe_down", dn == ref::Q2_0 && !g_dn_prep ? 142 : dn, b)); + K.cgu = a; + K.cdn = b; + if (std::getenv("PIPE_INFO")) + std::printf("pipelines: gu %s max threads %d tg mem %d | dn %s max threads %d tg mem %d\n", K.gu.name.c_str(), K.gu.max_threads, + K.gu.static_tg_mem, K.dn.name.c_str(), K.dn.max_threads, K.dn.static_tg_mem); + return K; + }; + auto args_for = [&](const Kernels& K, int nt, int n_expert) { + MoeArgs a{}; + a.gate_off = 0; + a.up_off = L.up_off; + a.down_off = L.down_off; + a.gu_stride = L.blob; + a.d_stride = L.blob; + a.gu_row = (uint32_t)L.gu_row; + a.d_row = (uint32_t)L.d_row; + a.n_embd = N_EMBD; + a.n_ff = N_FF; + a.n_tok = nt; + a.topk = TOPK; + a.n_expert = n_expert; + a.gx_gu = N_FF / rows_per_tg(K.cgu); + a.gx_dn = N_EMBD / rows_per_tg(K.cdn); + return a; + }; + + mtl::Buffer x = cx.alloc(4 * 8 * N_EMBD), h = cx.alloc(4 * 8 * TOPK * N_FF), ye = cx.alloc(4 * 8 * TOPK * N_EMBD), + y = cx.alloc(4 * 8 * N_EMBD); + mtl::Buffer XPX = cx.alloc(4 * (4 * PLANE_X + PLANE_X / 16)), XPH = cx.alloc(4 * (4 * PLANE_H + PLANE_H / 16)); + g_xpx = &XPX; + g_xph = &XPH; + g_gu_type = gu; + g_dn_type = dn; + ref::Rng rng(12345 + (seed - 1)); + for (int i = 0; i < 8 * N_EMBD; ++i) x.as()[i] = rng.sym(); + if (g_had_h) { // H (s * x) per token on the GPU, checked against the CPU's; x then holds the rotated rows + std::vector want(x.as(), x.as() + 8 * N_EMBD); + for (int t = 0; t < 8; ++t) strata::kernels::cpu::hadamard_rotate(want.data() + t * N_EMBD, N_EMBD, 128, sgx); + HadArgs ha{(uint32_t)N_EMBD, (uint32_t)N_EMBD, (uint32_t)N_EMBD, 1.0f / std::sqrt(128.0f), sgx ? 1u : 0u, 0u}; + mtl::Dispatch d; + d.p = &had32; + d.args = {mtl::Arg::B(x), mtl::Arg::B(x), mtl::Arg::B(sgx ? hsx : x), mtl::Arg::V(ha)}; + d.grid[0] = N_EMBD / 128; d.grid[1] = 8; d.tg[0] = 128; + cx.run({d}); + int bad = 0; + for (int i = 0; i < 8 * N_EMBD; ++i) bad += std::memcmp(&want[i], x.as() + i, 4) != 0; + std::printf("x rotation on the GPU vs the CPU: %d of %d values differ\n", bad, 8 * N_EMBD); + } + { + mtl::Pipeline pp = cx.pipeline(lib, "xprep"); + cx.run({prep_dispatch(pp, x, XPX, N_EMBD, 8)}); + } + + if (!bench) { + // ---------------- parity ---------------- + const int NE = tg_ && experts_first ? std::max(experts_first, 2 * TOPK) : 64; + mtl::Buffer W = cx.alloc(L.blob * NE); + std::vector pick; + for (int e = 0; e < NE; ++e) { + uint8_t* b = W.as() + e * L.blob; + if (tg_) { // real experts of the file: expert `ex` of the layer's three tensors + int ex; + if (experts_first) ex = e % experts_first; // (a sparse copy holds only the first experts: reused when NE is larger) + else do ex = (int)(rng.next() % 512); while (std::find(pick.begin(), pick.end(), ex) != pick.end()); + pick.push_back(ex); + std::memcpy(b, (const uint8_t*)gf.map + tg_->offset + (size_t)ex * L.gu_row * N_FF, L.gu_row * N_FF); + std::memcpy(b + L.up_off, (const uint8_t*)gf.map + tu_->offset + (size_t)ex * L.gu_row * N_FF, L.gu_row * N_FF); + std::memcpy(b + L.down_off, (const uint8_t*)gf.map + td_->offset + (size_t)ex * L.d_row * N_EMBD, L.d_row * N_EMBD); + } else { + ref::random_blocks(gu, b, 2 * N_FF, N_EMBD, 1000 + e + 100 * (seed - 1)); + ref::random_blocks(dn, b + L.down_off, N_EMBD, N_FF, 5000 + e + 100 * (seed - 1)); + } + } + std::vector ids; + std::vector w; + random_routing(rng, ntok, NE, ids, w, 0.3); + mtl::Buffer bid = cx.alloc(4 * ids.size()), bw = cx.alloc(4 * w.size()), plan = cx.alloc(4 * P_INTS); + std::memcpy(bid.ptr, ids.data(), 4 * ids.size()); + std::memcpy(bw.ptr, w.data(), 4 * w.size()); + Kernels K = make(cgu, cdn); + MoeArgs a = args_for(K, ntok, 512); + std::vector ds; + encode_layer(ds, K, W, 0, x, h, ye, y, bid, bw, plan, 0, a, true, true, true, true); + cx.run(ds); + // the plan + const int* P = plan.as(); + std::set distinct(ids.begin(), ids.end()); + bool plan_ok = P[P_NG] == (int)distinct.size() && P[P_IGU + 1] == P[P_NG] && P[P_IDN + 1] == P[P_NG]; + { + int gi = 0; + for (int e : distinct) { + if (P[P_GEXP + gi] != e) plan_ok = false; + int j = P[P_GSTART + gi]; + for (int i = 0; i < ntok * TOPK; ++i) + if (ids[i] == e) { + if (P[P_ETOK + j] != i / TOPK || P[P_TOKENT + i] != j) plan_ok = false; + float fw; + std::memcpy(&fw, &P[P_EW + j], 4); + if (fw != w[i]) plan_ok = false; + ++j; + } + if (j != P[P_GSTART + gi + 1]) plan_ok = false; + ++gi; + } + } + // the reference + Err eh, ey; + std::vector wg(N_FF * N_EMBD), wu(N_FF * N_EMBD), wd(N_EMBD * N_FF); + std::vector yref(ntok * N_EMBD, 0.0); + for (int i = 0; i < ntok * TOPK; ++i) { + const int e = ids[i], t = i / TOPK; + const uint8_t* b = W.as() + e * L.blob; + ref::dequant(gu, b, wg.data(), wg.size()); + ref::dequant(gu, b + L.up_off, wu.data(), wu.size()); + ref::dequant(dn, b + L.down_off, wd.data(), wd.size()); + const float* xt = x.as() + t * N_EMBD; + std::vector hr(N_FF); + const int ent = P[P_TOKENT + i]; + for (int r = 0; r < N_FF; ++r) { + double g = 0, u = 0; + for (int k = 0; k < N_EMBD; ++k) { + g += (double)wg[r * N_EMBD + k] * xt[k]; + u += (double)wu[r * N_EMBD + k] * xt[k]; + } + hr[r] = g / (1.0 + std::exp(-g)) * u; + } + if (g_had_h) { // the reference of the down input: H h, in float as the engine's CPU path does + std::vector hf(N_FF); + for (int r = 0; r < N_FF; ++r) hf[r] = (float)hr[r]; + strata::kernels::cpu::hadamard_rotate(hf.data(), N_FF, 128, sgh); + for (int r = 0; r < N_FF; ++r) hr[r] = hf[r]; + } + for (int r = 0; r < N_FF; ++r) eh.add(h.as()[ent * N_FF + r], hr[r]); + for (int r = 0; r < N_EMBD; ++r) { + double s = 0; + // the reference down uses the GPU's h (float) so that the y check isolates the down kernel ... + for (int k = 0; k < N_FF; ++k) s += (double)wd[r * N_FF + k] * h.as()[ent * N_FF + k]; + yref[t * N_EMBD + r] += w[i] * s; + } + } + for (int i = 0; i < ntok * N_EMBD; ++i) ey.add(y.as()[i], yref[i]); + const bool ok = plan_ok && eh.rel_l2() < 1e-5 && ey.rel_l2() < 1e-5 && eh.rel_max() < 1e-4 && ey.rel_max() < 1e-4; + std::printf("parity %s/%s ntok %d gu %d,%d,%d dn %d,%d,%d: plan %s (%d groups), h rel L2 %.2e max %.2e, " + "y rel L2 %.2e max %.2e %s\n", + ref::type_name(gu), ref::type_name(dn), ntok, cgu.L, cgu.R, cgu.S, cdn.L, cdn.R, cdn.S, + plan_ok ? "ok" : "WRONG", P[P_NG], eh.rel_l2(), eh.rel_max(), ey.rel_l2(), ey.rel_max(), + ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; + } + + // ---------------- bench ---------------- + const int NE = 512; + const size_t layer_bytes = L.blob * NE; + std::printf("arena: %d layers x %d experts = %.2f GB (max buffer %.1f GB)\n", layers, NE, layers * layer_bytes / 1e9, + cx.max_buffer_length() / 1e9); + mtl::Buffer W = cx.alloc(layer_bytes * layers); + { + // 64 distinct random experts, copied over the arena + std::vector tmp(L.blob * 64); + for (int e = 0; e < 64; ++e) { + ref::random_blocks(gu, tmp.data() + e * L.blob, 2 * N_FF, N_EMBD, 1000 + e); + ref::random_blocks(dn, tmp.data() + e * L.blob + L.down_off, N_EMBD, N_FF, 5000 + e); + } + for (size_t e = 0; e < (size_t)NE * layers; ++e) + std::memcpy(W.as() + e * L.blob, tmp.data() + (e % 64) * L.blob, L.blob); + } + std::vector> rids(reps); + std::vector> rw(reps); + mtl::Buffer bid = cx.alloc(4 * 8 * TOPK * reps), bw = cx.alloc(4 * 8 * TOPK * reps), plan = cx.alloc(4 * P_INTS * reps); + auto set_routing = [&](int nt) { + double groups = 0; + for (int r = 0; r < reps; ++r) { + random_routing(rng, nt, NE, rids[r], rw[r], share); + std::memcpy(bid.as() + r * 8 * TOPK, rids[r].data(), 4 * rids[r].size()); + std::memcpy(bw.as() + r * 8 * TOPK, rw[r].data(), 4 * rw[r].size()); + groups += std::set(rids[r].begin(), rids[r].end()).size(); + } + return groups / reps; + }; + auto run_cfg = [&](const Kernels& K, int nt, double groups, bool print) { + const MoeArgs a = args_for(K, nt, NE); + std::vector dp, dg, dd, dall; + for (int r = 0; r < reps; ++r) { + const size_t woff = (size_t)(r % layers) * layer_bytes; + std::vector one; + encode_layer(one, K, W, woff, x, h, ye, y, bid, bw, plan, (size_t)r * P_INTS * 4, a, true, false, false, false); + one[0].args[0] = mtl::Arg::B(bid, (size_t)r * 8 * TOPK * 4); + one[0].args[1] = mtl::Arg::B(bw, (size_t)r * 8 * TOPK * 4); + dp.push_back(one[0]); + std::vector g, d; + encode_layer(g, K, W, woff, x, h, ye, y, bid, bw, plan, (size_t)r * P_INTS * 4, a, false, true, false, false); + encode_layer(d, K, W, woff, x, h, ye, y, bid, bw, plan, (size_t)r * P_INTS * 4, a, false, false, true, false); + dg.push_back(g[0]); + for (auto& z : d) dd.push_back(z); + dall.push_back(one[0]); + dall.push_back(g[0]); + for (auto& z : d) dall.push_back(z); + } + cx.run(dp); // plans in place + cx.warm(dall, 0.1); + std::vector tp, tg, td, ta; + for (int it = 0; it < 7; ++it) { + tp.push_back(cx.run(dp) / reps); + tg.push_back(cx.run(dg) / reps); + td.push_back(cx.run(dd) / reps); + ta.push_back(cx.run(dall) / reps); + } + const double mp = median(tp), mg = median(tg), md = median(td), ma = median(ta); + const double bg = groups * 2.0 * N_FF * L.gu_row, bd = groups * N_EMBD * L.d_row; + if (print) + std::printf("ntok %d groups %5.1f | gu %2d,%d,%d %7.1f us %6.1f GB/s | dn %2d,%d,%d %7.1f us %6.1f GB/s | " + "plan %5.1f us | layer %7.1f us = %6.1f GB/s (%4.1f%% of 585)\n", + nt, groups, K.cgu.L, K.cgu.R, K.cgu.S, mg * 1e6, bg / mg / 1e9, K.cdn.L, K.cdn.R, K.cdn.S, + md * 1e6, bd / md / 1e9, mp * 1e6, ma * 1e6, (bg + bd) / ma / 1e9, + (bg + bd) / ma / 1e9 / 585 * 100); + return std::make_pair(mg, md); + }; + const std::vector gus = {{32, 1, 4}, {32, 2, 4}, {32, 4, 4}, {32, 2, 2}, {32, 2, 8}, {32, 1, 8}, + {16, 1, 4}, {16, 2, 4}, {8, 1, 4}, {8, 2, 4}, {8, 1, 2}}; + const std::vector dns = {{8, 1, 4}, {8, 2, 4}, {8, 4, 4}, {8, 1, 8}, {8, 2, 8}, {4, 1, 4}, + {4, 2, 4}, {4, 4, 4}, {16, 1, 4}, {16, 2, 4}, {32, 1, 4}, {2, 2, 4}}; + std::vector ntoks = {ntok}; + if (sweep) ntoks = {1, 2, 3, 4, 8}; + for (int nt : ntoks) { + const double groups = set_routing(nt); + if (sweep) { + // gate/up sweep with the default down, then down sweep with the best gate/up + Cfg bgu = cgu, bdn = cdn; + double best = 1e9; + for (Cfg c : gus) { + if (N_FF % rows_per_tg(c)) continue; + auto t = run_cfg(make(c, cdn), nt, groups, false); + std::printf(" gu %2d,%d,%d: %7.1f us %6.1f GB/s\n", c.L, c.R, c.S, t.first * 1e6, + groups * 2.0 * N_FF * L.gu_row / t.first / 1e9); + if (t.first < best) { best = t.first; bgu = c; } + } + best = 1e9; + for (Cfg c : dns) { + if (N_EMBD % rows_per_tg(c)) continue; + auto t = run_cfg(make(bgu, c), nt, groups, false); + std::printf(" dn %2d,%d,%d: %7.1f us %6.1f GB/s\n", c.L, c.R, c.S, t.second * 1e6, + groups * N_EMBD * L.d_row / t.second / 1e9); + if (t.second < best) { best = t.second; bdn = c; } + } + run_cfg(make(bgu, bdn), nt, groups, true); + } else { + run_cfg(make(cgu, cdn), nt, groups, true); + } + } + return 0; +} diff --git a/edge0-strata/metal/tests/fuse_parity.cpp b/edge0-strata/metal/tests/fuse_parity.cpp new file mode 100644 index 00000000..e67d9e58 --- /dev/null +++ b/edge0-strata/metal/tests/fuse_parity.cpp @@ -0,0 +1,160 @@ +// metal/tests/fuse_parity.cpp - the decode latency fusions against the kernels they replace, BITWISE, for 1..4 +// tokens: +// gemv_cat (gemv.metal) ssm_alpha + ssm_beta (BF16 2560 -> 48 + 48) in one dispatch == two gemv dispatches +// router_route the router GEMV + moe_route in one dispatch == gemv + moe_route (logits, ids, w, plan; +// twice in a row: the counter resets) +// gr_norm_comb (gr2.metal) the MoE combine folded into the next gr_norm == moe_combine_sh + gr_norm (R', xn) +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include + +namespace { +struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; +struct MoeArgs { + uint64_t gate_off, up_off, down_off, gu_stride, d_stride; + uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad; +}; +constexpr int P_INTS = 16 + 128 + 128 + 4 + 128 * 3, P_TOKENT = 16 + 128 + 128 + 4 + 128 * 2; +struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; +struct GlueArgs { uint32_t n, n_tok, n_embd, topk; }; +} // namespace + +int main() { + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "router.metal", "gr.metal", + "gr2.metal", "glue.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); + int fails = 0; + const int N = 2560, NO = 48, MAXT = 8; + const size_t rb = ref::row_bytes(ref::BF16, N); + mtl::Buffer Wa = cx.alloc(rb * NO), Wb = cx.alloc(rb * NO), xp = cx.alloc(4 * ((size_t)4 * MAXT * N + MAXT * N / 16)), + y1 = cx.alloc(4 * MAXT * NO), y2 = cx.alloc(4 * MAXT * NO), z1 = cx.alloc(4 * MAXT * NO), z2 = cx.alloc(4 * MAXT * NO); + ref::random_blocks(ref::BF16, Wa.as(), NO, N, 1, 0.05f); + ref::random_blocks(ref::BF16, Wb.as(), NO, N, 2, 0.05f); + ref::Rng r(3); + for (size_t i = 0; i < xp.bytes / 4; ++i) xp.as()[i] = r.sym(); + for (int nt = 1; nt <= 4; ++nt) { + mtl::Pipeline pg = cx.pipeline(lib, "gemv_30_32_1_4_1", mtl::Constants().i(0, nt)), + pc = cx.pipeline(lib, "gemvc_30_32_1_4", mtl::Constants().i(0, nt)); + const size_t s16 = (size_t)4 * 4 * MAXT * N; + GemvArgs a{(uint32_t)N, (uint32_t)NO, (uint32_t)nt, (uint32_t)rb, (uint32_t)N, (uint32_t)NO}; + auto g = [&](const mtl::Buffer& W, const mtl::Buffer& y) { + mtl::Dispatch d; + d.p = &pg; + d.args = {mtl::Arg::B(W), mtl::Arg::B(xp), mtl::Arg::B(xp, s16), mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = NO / 4; + d.tg[0] = 128; + return d; + }; + cx.run({g(Wa, y1), g(Wb, y2)}); + GemvArgs ac = a; + ac.n_out = 2 * NO; + const uint32_t sp = NO; + mtl::Dispatch d; + d.p = &pc; + d.args = {mtl::Arg::B(Wa), mtl::Arg::B(xp), mtl::Arg::B(xp, s16), mtl::Arg::B(z1), mtl::Arg::V(ac), mtl::Arg::B(Wb), + mtl::Arg::B(z2), mtl::Arg::V(sp)}; + d.grid[0] = 2 * NO / 4; + d.tg[0] = 128; + cx.run({d}); + const bool ok = std::memcmp(y1.ptr, z1.ptr, 4 * nt * NO) == 0 && std::memcmp(y2.ptr, z2.ptr, 4 * nt * NO) == 0; + fails += !ok; + std::printf("gemv_cat (alpha + beta) ntok %d: bitwise %s\n", nt, ok ? "PASS" : "FAIL"); + } + // ---- router_route + { + for (int rty : {ref::BF16, ref::F32}) { + const int NE = 512; + const size_t rbr = ref::row_bytes(rty, N); + mtl::Buffer Wr = cx.alloc(rbr * NE), lg1 = cx.alloc(4 * MAXT * NE), lg2 = cx.alloc(4 * MAXT * NE), + id1 = cx.alloc(4 * MAXT * 10), id2 = cx.alloc(4 * MAXT * 10), w1 = cx.alloc(4 * MAXT * 10), + w2 = cx.alloc(4 * MAXT * 10), pl1 = cx.alloc(4 * P_INTS), pl2 = cx.alloc(4 * P_INTS), cnt = cx.alloc(16); + ref::random_blocks(rty, Wr.as(), NE, N, 9, 0.05f); + mtl::Pipeline pr = cx.pipeline(lib, "moe_route"); + for (int nt = 1; nt <= 4; ++nt) { + mtl::Pipeline pg = cx.pipeline(lib, "gemv_" + std::to_string(rty) + "_32_1_4_1", mtl::Constants().i(0, nt)), + pf = cx.pipeline(lib, "router_route_" + std::to_string(rty), mtl::Constants().i(0, nt)); + const size_t s16 = (size_t)4 * 4 * MAXT * N; + GemvArgs ga{(uint32_t)N, (uint32_t)NE, (uint32_t)nt, (uint32_t)rbr, (uint32_t)N, (uint32_t)NE}; + MoeArgs ma{}; + ma.n_tok = nt; + ma.topk = 10; + ma.n_expert = NE; + ma.gx_gu = 5; + ma.gx_dn = 20; + mtl::Dispatch g; + g.p = &pg; + g.args = {mtl::Arg::B(Wr), mtl::Arg::B(xp), mtl::Arg::B(xp, s16), mtl::Arg::B(lg1), mtl::Arg::V(ga)}; + g.grid[0] = NE / 4; + g.tg[0] = 128; + mtl::Dispatch rt; + rt.p = ≺ + rt.args = {mtl::Arg::B(lg1), mtl::Arg::B(id1), mtl::Arg::B(w1), mtl::Arg::B(pl1), mtl::Arg::V(ma)}; + rt.tg[0] = 512; + std::memset(pl1.ptr, 0, pl1.bytes); + cx.run({g, rt}); + bool ok = true; + for (int rep_ = 0; rep_ < 2; ++rep_) { + std::memset(pl2.ptr, 0, pl2.bytes); + mtl::Dispatch f; + f.p = &pf; + f.args = {mtl::Arg::B(Wr), mtl::Arg::B(xp), mtl::Arg::B(xp, s16), mtl::Arg::B(lg2), mtl::Arg::V(ga), mtl::Arg::B(id2), + mtl::Arg::B(w2), mtl::Arg::B(pl2), mtl::Arg::V(ma), mtl::Arg::B(cnt)}; + f.grid[0] = NE / 16; + f.tg[0] = 512; + cx.run({f}); + ok = ok && std::memcmp(lg1.ptr, lg2.ptr, 4 * nt * NE) == 0 && std::memcmp(id1.ptr, id2.ptr, 4 * nt * 10) == 0 && + std::memcmp(w1.ptr, w2.ptr, 4 * nt * 10) == 0 && std::memcmp(pl1.ptr, pl2.ptr, 4 * P_INTS) == 0 && + cnt.as()[0] == 0; + } + fails += !ok; + std::printf("router_route %s ntok %d: logits, ids, w, plan bitwise (2 runs) %s\n", ref::type_name(rty), nt, + ok ? "PASS" : "FAIL"); + } + } + } + // ---- gr_norm_comb + { + const int HC = 4, TOPK = 10; + mtl::Buffer R0 = cx.alloc(4 * MAXT * HC * N), R1 = cx.alloc(4 * MAXT * HC * N), R2 = cx.alloc(4 * MAXT * HC * N), + ye = cx.alloc(4 * (size_t)MAXT * TOPK * N), ysh = cx.alloc(4 * MAXT * N), sgv = cx.alloc(4 * MAXT), + plan = cx.alloc(4 * P_INTS), inj = cx.alloc(4 * MAXT * HC), wn = cx.alloc(4 * HC * N), bo = cx.alloc(4 * MAXT * N), + x1 = cx.alloc(4 * MAXT * HC * N), x2 = cx.alloc(4 * MAXT * HC * N); + for (size_t i = 0; i < R0.bytes / 4; ++i) R0.as()[i] = r.sym(); + for (size_t i = 0; i < ye.bytes / 4; ++i) ye.as()[i] = r.sym(); + for (size_t i = 0; i < ysh.bytes / 4; ++i) ysh.as()[i] = r.sym(); + for (int i = 0; i < MAXT; ++i) sgv.as()[i] = r.sym() * 3; + for (int i = 0; i < MAXT * HC; ++i) inj.as()[i] = r.sym() * 3; + for (int i = 0; i < HC * N; ++i) wn.as()[i] = 0.5f + r.uni(); + for (int i = 0; i < MAXT * TOPK; ++i) plan.as()[P_TOKENT + i] = (i * 7) % (MAXT * TOPK); + mtl::Pipeline pcm = cx.pipeline(lib, "moe_combine_sh"), pn = cx.pipeline(lib, "gr_norm"), pnc = cx.pipeline(lib, "gr_norm_comb"); + for (int nt = 1; nt <= 4; ++nt) { + std::memcpy(R1.ptr, R0.ptr, R0.bytes); + std::memcpy(R2.ptr, R0.ptr, R0.bytes); + GlueArgs ga{(uint32_t)N, (uint32_t)nt, (uint32_t)N, TOPK}; + GrArgs gr{(uint32_t)N, HC, 320, (uint32_t)nt, 1, 1e-6f}; + mtl::Dispatch c; c.p = &pcm; + c.args = {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::B(bo), mtl::Arg::V(ga)}; + c.grid[0] = (nt * N + 255) / 256; c.tg[0] = 256; + mtl::Dispatch n1; n1.p = &pn; + n1.args = {mtl::Arg::B(R1), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x1), mtl::Arg::V(gr)}; + n1.grid[0] = HC; n1.grid[1] = nt; n1.tg[0] = 256; + cx.run({c, n1}); + const uint32_t tk = TOPK; + mtl::Dispatch n2; n2.p = &pnc; + n2.args = {mtl::Arg::B(R2), mtl::Arg::B(ye), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(x2), mtl::Arg::V(gr), + mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sgv), mtl::Arg::V(tk)}; + n2.grid[0] = HC; n2.grid[1] = nt; n2.tg[0] = 256; + cx.run({n2}); + const bool ok = std::memcmp(R1.ptr, R2.ptr, 4 * (size_t)nt * HC * N) == 0 && std::memcmp(x1.ptr, x2.ptr, 4 * (size_t)nt * HC * N) == 0; + fails += !ok; + std::printf("gr_norm_comb ntok %d: R', xn bitwise %s\n", nt, ok ? "PASS" : "FAIL"); + } + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/gdn_chunk_bench.cpp b/edge0-strata/metal/tests/gdn_chunk_bench.cpp new file mode 100644 index 00000000..c439c10a --- /dev/null +++ b/edge0-strata/metal/tests/gdn_chunk_bench.cpp @@ -0,0 +1,374 @@ +// metal/tests/gdn_chunk_bench.cpp - GDN prompt stage microbenchmark and numerical check (30 Oct GDN-M, phase 1). +// The sequential scans of gdn_prompt.metal (gdn_scanr_2/4/8, production = gdn_scanr_4) against the chunked C = 32 form +// (gdn_cprep -> gdn_cstate -> gdn_cout) at the real Flash-Next shape (h_k 16, h_v 48, S 128, 10240 conv channels), +// T = 512 .. 32768, with the same gdn_pre in front and gdn_post behind (the whole GDN stage). Checks against an FP64 +// sequential oracle on a few heads (output o and final state), the pipeline thread limits, and continuation (a call +// split in two, chunked or sequential on either side, the second call on buffer offsets). +// gdn_chunk_bench [--T 512,2048,..] [--check] [--regime normal|strong] [--heads 0,17,47] [--reps 5] [--noscan] +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int HK = 16, HV = 48, CH = 10240, S = 128, GC = 32, NSC = 14; +struct GdnPArgs { uint32_t n_tok, h_k, h_v, channels; float eps; uint32_t f32; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +std::vector ints(const char* s) { + std::vector r; + for (const char* p = s; *p;) { r.push_back(std::atoi(p)); while (*p && *p != ',') ++p; if (*p) ++p; } + return r; +} +} // namespace + +int main(int argc, char** argv) { + std::vector Ts = {512, 1024, 2048, 4096, 8192, 16384, 32768}, heads = {0, 17, 47}; + bool check = false, noscan = false, strong = false; + int reps = 5; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--T") Ts = ints(argv[++i]); + else if (a == "--heads") heads = ints(argv[++i]); + else if (a == "--check") check = true; + else if (a == "--noscan") noscan = true; + else if (a == "--regime") strong = std::string(argv[++i]) == "strong"; + else if (a == "--reps") reps = std::atoi(argv[++i]); + } + const int TMAX = *std::max_element(Ts.begin(), Ts.end()); + mtl::Context cx; + const int lib = cx.compile({"gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("device %s, regime %s, TMAX %d\n", cx.device_name().c_str(), strong ? "strong-decay" : "normal", TMAX); + auto F = [&](size_t n) { return cx.alloc(4 * n); }; + const size_t NCM = (TMAX + GC - 1) / GC; + mtl::Buffer qkv = F((size_t)TMAX * CH), al = F((size_t)TMAX * HV), be = F((size_t)TMAX * HV), z = F((size_t)TMAX * HV * S), + cw = F(CH * 4), dt = F(HV), sa = F(HV), gam = F(S), st0 = F((size_t)S * HV * S), stX = F((size_t)S * HV * S), + stRef = F((size_t)S * HV * S), h0 = F(CH * 3), hOut = F(CH * 3), yh = cx.alloc(2 * (size_t)TMAX * HV * S), + yB = F((size_t)TMAX * HV * S), qn = F((size_t)(TMAX + 32) * HK * S), kn = F((size_t)(TMAX + 32) * HK * S), + vc = F((size_t)TMAX * HV * S), gb = F(2 * (size_t)TMAX * HV), kq = F((size_t)TMAX * HK), + o = F((size_t)TMAX * HV * S), oRef = F((size_t)TMAX * HV * S), cU = F((size_t)HV * NCM * GC * S), + cW = F((size_t)HV * NCM * GC * S), clg = F((size_t)HV * NCM * GC), cD = F((size_t)HV * NCM * GC * S), cM = F((size_t)HV * NCM * GC * GC); + ref::Rng r(21); + for (size_t i = 0; i < (size_t)TMAX * CH; ++i) qkv.as()[i] = r.sym() * 2.0f; + for (size_t i = 0; i < (size_t)TMAX * HV; ++i) { al.as()[i] = r.sym() * 3.0f; be.as()[i] = r.sym() * 2.0f; } + for (size_t i = 0; i < (size_t)TMAX * HV * S; ++i) z.as()[i] = r.sym() * 2.0f; + for (int i = 0; i < CH * 4; ++i) cw.as()[i] = r.sym() * 0.6f; + for (int i = 0; i < HV; ++i) { dt.as()[i] = r.sym(); sa.as()[i] = strong ? -5.0f - r.uni() * 25.0f : -0.2f - r.uni() * 2.0f; } + for (int i = 0; i < S; ++i) gam.as()[i] = 0.5f + r.uni(); + for (int i = 0; i < S * HV * S; ++i) st0.as()[i] = r.sym() * 0.1f; + for (int i = 0; i < CH * 3; ++i) h0.as()[i] = r.sym(); + + mtl::Pipeline pre = cx.pipeline(lib, "gdn_pre"), post = cx.pipeline(lib, "gdn_post"); + // the existing sequential scans: {name, threadgroups, threads}; index 1 = gdn_scanr_4 = production + struct Sv { const char* nm; int grid, tg; }; + const Sv svs[NSC] = {{"gdn_scanr_2", HV * 4, 64}, {"gdn_scanr_4", HV * 4, 128}, {"gdn_scanr_8", HV * 4, 256}, + {"gdn_scanr_4_8", HV * 4, 128}, {"gdn_scan_v", HV * 4, 256}, {"gdn_scan_1", HV * 4, 32}, + {"gdn_scan_2", HV * 2, 64}, {"gdn_scan_4", HV, 128}, + {"gdn_scanr_4_4", HV * 4, 128}, {"gdn_scanr_4_2", HV * 4, 128}, {"gdn_scanr_8_8", HV * 4, 256}, + {"gdn_scanr_8_4", HV * 4, 256}, {"gdn_scanr_2_8", HV * 4, 64}, {"gdn_scanr_2_4", HV * 4, 64}}; + std::vector scan; + for (const Sv& v : svs) scan.push_back(cx.pipeline(lib, v.nm, std::string(v.nm) == "gdn_scan_v" || std::string(v.nm).find("scanr") != std::string::npos ? mtl::Constants().i(6, 0) : mtl::Constants())); + mtl::Pipeline pcp = cx.pipeline(lib, "gdn_cprep"), pcs = cx.pipeline(lib, "gdn_cstate"), pco = cx.pipeline(lib, "gdn_cout"), pcp2 = cx.pipeline(lib, "gdn_cprep2"), pcs2 = cx.pipeline(lib, "gdn_cstate2"); + bool v2 = false; // chunk_d builds the v2 chunked form (gdn_cprep2 + gdn_cstate2) + std::printf("pipeline limits (max threads / simd / static tg mem; threads asked):"); + for (size_t i = 0; i < scan.size(); ++i) std::printf(" %s %d/%d/%d (%d)", svs[i].nm, scan[i].max_threads, scan[i].simd_width, scan[i].static_tg_mem, svs[i].tg); + std::printf("\n pre %d/%d/%d (256) cprep %d/%d/%d (256) cstate %d/%d/%d (256) cout %d/%d/%d (128) post %d/%d/%d (128)\n", + pre.max_threads, pre.simd_width, pre.static_tg_mem, pcp.max_threads, pcp.simd_width, pcp.static_tg_mem, pcs.max_threads, + pcs.simd_width, pcs.static_tg_mem, pco.max_threads, pco.simd_width, pco.static_tg_mem, post.max_threads, post.simd_width, + post.static_tg_mem); + std::printf("scratch at T: U+W+delta+lg = %.1f KB per token (%.0f MB at T 2048, %.0f MB at T 4096)\n", + (3.0 * HV * S * 4 + HV * 4) / 1024.0, (3.0 * HV * S * 4 + HV * 4) * 2048 / 1e6, + (3.0 * HV * S * 4 + HV * 4) * 4096 / 1e6); + + // token range [t0, t0 + n) as buffer offsets + auto pre_d = [&](int t0, int n, std::vector& ds, bool first) { + GdnPArgs a{(uint32_t)n, HK, HV, CH, 1e-6f, 1u}; + mtl::Dispatch d; + d.p = ⪯ + // history: the chunk reads hist (h0 for the first call, hOut afterwards), writes hOut + d.args = {mtl::Arg::B(qkv, (size_t)t0 * CH * 4), mtl::Arg::B(first ? h0 : hOut), mtl::Arg::B(cw), + mtl::Arg::B(al, (size_t)t0 * HV * 4), mtl::Arg::B(be, (size_t)t0 * HV * 4), mtl::Arg::B(dt), mtl::Arg::B(sa), + mtl::Arg::B(qn, (size_t)t0 * HK * S * 4), mtl::Arg::B(kn, (size_t)t0 * HK * S * 4), + mtl::Arg::B(vc, (size_t)t0 * HV * S * 4), mtl::Arg::B(gb, (size_t)t0 * HV * 8), mtl::Arg::B(kq, (size_t)t0 * HK * 4), + mtl::Arg::B(hOut), mtl::Arg::V(a)}; + d.grid[0] = HK + HV / 2; d.grid[1] = n; d.tg[0] = 256; + ds.push_back(d); + }; + auto scan_d = [&](int which, int t0, int n, const mtl::Buffer& st, const mtl::Buffer& out, std::vector& ds) { + GdnPArgs a{(uint32_t)n, HK, HV, CH, 1e-6f, 1u}; + mtl::Dispatch d; + d.p = &scan[which]; + d.args = {mtl::Arg::B(st), mtl::Arg::B(qn, (size_t)t0 * HK * S * 4), mtl::Arg::B(kn, (size_t)t0 * HK * S * 4), + mtl::Arg::B(vc, (size_t)t0 * HV * S * 4), mtl::Arg::B(gb, (size_t)t0 * HV * 8), mtl::Arg::B(kq, (size_t)t0 * HK * 4), + mtl::Arg::B(out, (size_t)t0 * HV * S * 4), mtl::Arg::V(a)}; + d.grid[0] = svs[which].grid; d.tg[0] = svs[which].tg; + ds.push_back(d); + }; + auto chunk_d = [&](int t0, int n, const mtl::Buffer& st, const mtl::Buffer& out, std::vector& ds) { + GdnPArgs a{(uint32_t)n, HK, HV, CH, 1e-6f, 1u}; + const uint32_t nc = (n + GC - 1) / GC; + if (v2) { + mtl::Dispatch f1; f1.p = &pcp2; + f1.args = {mtl::Arg::B(qn, (size_t)t0 * HK * S * 4), mtl::Arg::B(kn, (size_t)t0 * HK * S * 4), + mtl::Arg::B(vc, (size_t)t0 * HV * S * 4), mtl::Arg::B(gb, (size_t)t0 * HV * 8), mtl::Arg::B(cU), mtl::Arg::B(cW), + mtl::Arg::B(clg), mtl::Arg::B(cM), mtl::Arg::V(a)}; + f1.grid[0] = HV; f1.grid[1] = nc; f1.tg[0] = 256; ds.push_back(f1); + mtl::Dispatch f2; f2.p = &pcs2; + f2.args = {mtl::Arg::B(st), mtl::Arg::B(qn, (size_t)t0 * HK * S * 4), mtl::Arg::B(kn, (size_t)t0 * HK * S * 4), + mtl::Arg::B(cU), mtl::Arg::B(cW), mtl::Arg::B(clg), mtl::Arg::B(cM), mtl::Arg::B(out, (size_t)t0 * HV * S * 4), + mtl::Arg::V(a)}; + f2.grid[0] = HV * 4; f2.tg[0] = 128; ds.push_back(f2); + return; + } + mtl::Dispatch e1; e1.p = &pcp; + e1.args = {mtl::Arg::B(qn, (size_t)t0 * HK * S * 4), mtl::Arg::B(kn, (size_t)t0 * HK * S * 4), + mtl::Arg::B(vc, (size_t)t0 * HV * S * 4), mtl::Arg::B(gb, (size_t)t0 * HV * 8), mtl::Arg::B(cU), mtl::Arg::B(cW), + mtl::Arg::B(clg), mtl::Arg::V(a)}; + e1.grid[0] = HV; e1.grid[1] = nc; e1.tg[0] = 256; ds.push_back(e1); + mtl::Dispatch e2; e2.p = &pcs; + e2.args = {mtl::Arg::B(st), mtl::Arg::B(qn, (size_t)t0 * HK * S * 4), mtl::Arg::B(kn, (size_t)t0 * HK * S * 4), + mtl::Arg::B(cU), mtl::Arg::B(cW), mtl::Arg::B(clg), mtl::Arg::B(cD), mtl::Arg::B(out, (size_t)t0 * HV * S * 4), + mtl::Arg::V(a)}; + e2.grid[0] = HV * 4; e2.tg[0] = 256; ds.push_back(e2); + mtl::Dispatch e3; e3.p = &pco; + e3.args = {mtl::Arg::B(qn, (size_t)t0 * HK * S * 4), mtl::Arg::B(kn, (size_t)t0 * HK * S * 4), mtl::Arg::B(cD), + mtl::Arg::B(clg), mtl::Arg::B(out, (size_t)t0 * HV * S * 4), mtl::Arg::V(a)}; + e3.grid[0] = HV; e3.grid[1] = nc; e3.tg[0] = 128; ds.push_back(e3); + }; + auto post_d = [&](int t0, int n, const mtl::Buffer& out, std::vector& ds) { + GdnPArgs a{(uint32_t)n, HK, HV, CH, 1e-6f, 1u}; + mtl::Dispatch d; + d.p = &post; + d.args = {mtl::Arg::B(out, (size_t)t0 * HV * S * 4), mtl::Arg::B(z, (size_t)t0 * HV * S * 4), mtl::Arg::B(gam), + mtl::Arg::B(yh, (size_t)t0 * HV * S * 2), mtl::Arg::B(yB, (size_t)t0 * HV * S * 4), mtl::Arg::V(a)}; + d.grid[0] = HV; d.grid[1] = n; d.tg[0] = 128; + ds.push_back(d); + }; + mtl::Pipeline pf[3] = {cx.pipeline(lib, "gdn_scanf_4"), cx.pipeline(lib, "gdn_scanf_8"), cx.pipeline(lib, "gdn_scanf_2")}; + std::printf("fused scan+post pipelines (max threads / static tg mem; 512 asked): TB4 %d/%d TB8 %d/%d TB2 %d/%d\n", pf[0].max_threads, + pf[0].static_tg_mem, pf[1].max_threads, pf[1].static_tg_mem, pf[2].max_threads, pf[2].static_tg_mem); + auto fused_d = [&](int which, int t0, int n, const mtl::Buffer& st, std::vector& ds) { + GdnPArgs a{(uint32_t)n, HK, HV, CH, 1e-6f, 1u}; + mtl::Dispatch d; + d.p = &pf[which]; + d.args = {mtl::Arg::B(st), mtl::Arg::B(qn, (size_t)t0 * HK * S * 4), mtl::Arg::B(kn, (size_t)t0 * HK * S * 4), + mtl::Arg::B(vc, (size_t)t0 * HV * S * 4), mtl::Arg::B(gb, (size_t)t0 * HV * 8), mtl::Arg::B(kq, (size_t)t0 * HK * 4), + mtl::Arg::B(z, (size_t)t0 * HV * S * 4), mtl::Arg::B(gam), mtl::Arg::B(yh, (size_t)t0 * HV * S * 2), + mtl::Arg::B(yB, (size_t)t0 * HV * S * 4), mtl::Arg::V(a)}; + d.grid[0] = HV; d.tg[0] = 512; + ds.push_back(d); + }; + auto timeit = [&](const std::vector& ds) { + cx.warm(ds, 0.3); + std::vector t; + for (int i = 0; i < reps; ++i) t.push_back(cx.run(ds)); + return median(t) * 1e3; + }; + auto copy_st0 = [&](mtl::Buffer& dst) { std::memcpy(dst.ptr, st0.ptr, (size_t)S * HV * S * 4); }; + + // FP64 oracle for head h over tokens [t0, t0 + n) from the prepared tensors; state in/out [S][S] (i, j) + auto oracle = [&](int h, int t0, int n, std::vector& St, std::vector& out) { + const int hk = h % HK; + out.assign((size_t)n * S, 0.0); + std::vector a(S), b(S); + for (int t = t0; t < t0 + n; ++t) { + const float* k = kn.as() + ((size_t)t * HK + hk) * S; + const float* q = qn.as() + ((size_t)t * HK + hk) * S; + const float* v = vc.as() + ((size_t)t * HV + h) * S; + const double g = gb.as()[((size_t)t * HV + h) * 2], bt = gb.as()[((size_t)t * HV + h) * 2 + 1]; + double kqv = 0; + for (int i = 0; i < S; ++i) kqv += (double)k[i] * q[i]; + std::fill(a.begin(), a.end(), 0.0); std::fill(b.begin(), b.end(), 0.0); + for (int i = 0; i < S; ++i) { + const double* row = &St[(size_t)i * S]; + const double ki = k[i], qi = q[i]; + for (int j = 0; j < S; ++j) { a[j] += row[j] * ki; b[j] += row[j] * qi; } + } + std::vector dl(S); + for (int j = 0; j < S; ++j) { + dl[j] = (v[j] - g * a[j]) * bt; + out[(size_t)(t - t0) * S + j] = (g * b[j] + dl[j] * kqv) / std::sqrt((double)S); + } + for (int i = 0; i < S; ++i) { + double* row = &St[(size_t)i * S]; + for (int j = 0; j < S; ++j) row[j] = g * row[j] + (double)k[i] * dl[j]; + } + } + }; + auto initial = [&](int h, std::vector& St) { + St.resize((size_t)S * S); + for (int i = 0; i < S; ++i) for (int j = 0; j < S; ++j) St[(size_t)i * S + j] = st0.as()[((size_t)i * HV + h) * S + j]; + }; + struct Err { double l2, mx; int nonfinite; }; + auto cmp_o = [&](const mtl::Buffer& ob, int h, int t0, int n, const std::vector& ro) { + double se = 0, sr = 0, mx = 0, rmax = 0; int nf = 0; + for (int t = 0; t < n; ++t) for (int j = 0; j < S; ++j) { + const float x = ob.as()[((size_t)(t0 + t) * HV + h) * S + j]; + if (!std::isfinite(x)) { ++nf; continue; } + const double e = x - ro[(size_t)t * S + j]; + se += e * e; sr += ro[(size_t)t * S + j] * ro[(size_t)t * S + j]; mx = std::max(mx, std::fabs(e)); rmax = std::max(rmax, std::fabs(ro[(size_t)t * S + j])); + } + return Err{std::sqrt(se / std::max(sr, 1e-30)), mx / std::max(rmax, 1e-30), nf}; + }; + auto cmp_s = [&](const mtl::Buffer& sb, int h, const std::vector& rs) { + double se = 0, sr = 0, mx = 0, rmax = 0; int nf = 0; + for (int i = 0; i < S; ++i) for (int j = 0; j < S; ++j) { + const float x = sb.as()[((size_t)i * HV + h) * S + j]; + if (!std::isfinite(x)) { ++nf; continue; } + const double e = x - rs[(size_t)i * S + j]; + se += e * e; sr += rs[(size_t)i * S + j] * rs[(size_t)i * S + j]; mx = std::max(mx, std::fabs(e)); rmax = std::max(rmax, std::fabs(rs[(size_t)i * S + j])); + } + return Err{std::sqrt(se / std::max(sr, 1e-30)), mx / std::max(rmax, 1e-30), nf}; + }; + + // prepared tensors once for all TMAX tokens (gdn_pre per call carries the history; here one call, the tokens independent + // of the scan variant). For the continuation check the pre runs per split below. + { + std::vector ds; + // gdn_pre on a very long T reads hist only for the first 3 tokens; one call over TMAX is the same as pieces + pre_d(0, TMAX, ds, true); + cx.run(ds); + } + int fails = 0; + if (check) { + std::printf("== numerical checks vs FP64 oracle (rel L2 / max-abs over max|ref|, per head)\n"); + for (int T : Ts) { + // the variants, each from st0, over tokens [0, T) + struct V { const char* name; int kind; }; // kind: 0..2 scanr IB 2/4/8, 10 chunked + for (V v : {V{"scanr_4", 1}, V{"chunked", 10}, V{"chunked2", 11}}) { + copy_st0(stX); + std::vector ds; + v2 = v.kind == 11; if (v.kind >= 10) chunk_d(0, T, stX, o, ds); else scan_d(v.kind, 0, T, stX, o, ds); v2 = false; + cx.run(ds); + for (int h : heads) { + std::vector St, ro; initial(h, St); oracle(h, 0, T, St, ro); + Err eo = cmp_o(o, h, 0, T, ro), es = cmp_s(stX, h, St); + const bool ok = eo.l2 < 3e-3 && es.l2 < 3e-3 && !eo.nonfinite && !es.nonfinite; + fails += !ok; + std::printf("T %5d %-8s head %2d: o rel %.2e (max %.2e) state rel %.2e (max %.2e) nonfinite %d %s\n", T, v.name, + h, eo.l2, eo.mx, es.l2, es.mx, eo.nonfinite + es.nonfinite, ok ? "PASS" : "FAIL"); + if (T > 8192) break; // one head at the largest sizes (oracle cost) + } + } + } + // continuation: T = A + B, second call on offsets and the first call's state; every pairing + std::printf("== continuation (state carried across two calls, 2nd call on buffer offsets) vs oracle over the whole range\n"); + const std::pair splits[] = {{37, 963}, {1024, 1024}, {2048, 2047}, {31, 33}, {1, 4095}}; + for (auto [A, B] : splits) { + if (A + B > TMAX) continue; + const char* nm[8] = {"chunk+chunk", "scan+chunk", "chunk+scan", "scan+scan", "chunk2+chunk2", "scan+chunk2", "chunk2+scan", "-"}; + for (int m = 0; m < 8; ++m) { + v2 = m >= 4; + if (v2 && m == 7) continue; // scan+scan has no chunked part + copy_st0(stX); + std::vector d1, d2; + if (m % 4 == 0 || m % 4 == 2) chunk_d(0, A, stX, o, d1); else scan_d(1, 0, A, stX, o, d1); + if (m % 4 == 0 || m % 4 == 1) chunk_d(A, B, stX, o, d2); else scan_d(1, A, B, stX, o, d2); + cx.run(d1); cx.run(d2); v2 = false; + for (int h : {heads[0], heads.back()}) { + std::vector St, ro; initial(h, St); oracle(h, 0, A + B, St, ro); + Err eo = cmp_o(o, h, 0, A + B, ro), es = cmp_s(stX, h, St); + const bool ok = eo.l2 < 3e-3 && es.l2 < 3e-3 && !eo.nonfinite && !es.nonfinite; + fails += !ok; + std::printf("split %5d + %5d %-14s head %2d: o rel %.2e state rel %.2e %s\n", A, B, nm[m], h, eo.l2, es.l2, ok ? "PASS" : "FAIL"); + } + } + } + std::printf("== bitwise: gdn_scanr_4_8 / _4_4 / _4_2 / _8_8 / _2_8 o and state vs gdn_scanr_4\n"); + for (int T : Ts) { + if (T > 16384) continue; + copy_st0(stX); + std::vector d0; scan_d(1, 0, T, stX, o, d0); cx.run(d0); + std::vector oref(o.as(), o.as() + (size_t)T * HV * S), sref(stX.as(), stX.as() + (size_t)S * HV * S); + for (int w : {3, 8, 9, 10, 12}) { + copy_st0(stX); + std::memset(o.ptr, 0, (size_t)T * HV * S * 4); + std::vector d1; scan_d(w, 0, T, stX, o, d1); cx.run(d1); + const bool eq = std::memcmp(o.ptr, oref.data(), oref.size() * 4) == 0 && std::memcmp(stX.ptr, sref.data(), sref.size() * 4) == 0; + fails += !eq; + std::printf("T %5d %-14s %s\n", T, svs[w].nm, eq ? "BITWISE EQUAL" : "DIFFERS"); + } + } + std::printf("== fused scan+post vs pre/scanr_4/post: y rel L2, fp16 y mismatches, state rel L2\n"); + for (int T : Ts) { + if (T > 8192) continue; + copy_st0(stX); + std::vector d0; scan_d(1, 0, T, stX, o, d0); post_d(0, T, o, d0); cx.run(d0); + std::vector yref(yB.as(), yB.as() + (size_t)T * HV * S), sref(stX.as(), stX.as() + (size_t)S * HV * S); + std::vector href(yh.as(), yh.as() + (size_t)T * HV * S); + for (int w = 0; w < 3; ++w) { + copy_st0(stX); + std::memset(yB.ptr, 0, (size_t)T * HV * S * 4); + std::vector d1; fused_d(w, 0, T, stX, d1); cx.run(d1); + double se = 0, sr = 0, ss = 0, ssr = 0; size_t hm = 0; + for (size_t i = 0; i < (size_t)T * HV * S; ++i) { const double e = yB.as()[i] - yref[i]; se += e * e; sr += (double)yref[i] * yref[i]; hm += yh.as()[i] != href[i]; } + for (size_t i = 0; i < (size_t)S * HV * S; ++i) { const double e = stX.as()[i] - sref[i]; ss += e * e; ssr += (double)sref[i] * sref[i]; } + const bool ok = std::sqrt(se / sr) < 1e-5 && std::sqrt(ss / ssr) < 1e-5; + fails += !ok; + std::printf("T %5d fused TB%-2d: y rel %.2e, fp16 mismatches %zu, state rel %.2e %s\n", T, w == 0 ? 4 : w == 1 ? 8 : 2, std::sqrt(se / sr), hm, std::sqrt(ss / ssr), ok ? "PASS" : "FAIL"); + } + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + } + // timing: gdn_pre, the scan variants, chunked parts, gdn_post; the stage = pre + scan + post + std::printf("== timing (GPU ms per call, median of %d; stage = pre + scan + post)\n", reps); + std::printf("%6s | %7s %7s |", "T", "pre", "post"); + for (const Sv& v : svs) std::printf(" %s", v.nm + 4); + std::printf(" | chunked (cprep cstate cout) | stage scanr_4 / chunked / x\n"); + for (int T : Ts) { + std::vector dp, dpo, dc; + pre_d(0, T, dp, true); + post_d(0, T, o, dpo); + const double tp = timeit(dp), tpo = timeit(dpo); + std::printf("%6d | %7.3f %7.3f |", T, tp, tpo); + double ts1 = 0, tbest = 1e30; + int wbest = 1; + for (int w = 0; w < NSC; ++w) { + std::vector d1; + double t = 0; + if (!noscan) { scan_d(w, 0, T, stX, o, d1); t = timeit(d1); } + if (w == 1) ts1 = t; + if (t > 0 && t < tbest) { tbest = t; wbest = w; } + std::printf(" %8.3f", t); + } + chunk_d(0, T, stX, o, dc); + const double tc = timeit(dc), tc1 = timeit({dc[0]}), tc2 = timeit({dc[1]}), tc3 = timeit({dc[2]}); + const double base = tp + ts1 + tpo, ch = tp + tc + tpo; + std::printf(" | %8.3f (%7.3f %7.3f %7.3f) | %8.3f %8.3f %5.2f | best %s stage %.3f = x%.3f (scan x%.3f)\n", tc, tc1, tc2, tc3, base, ch, base / ch, svs[wbest].nm, tp + tbest + tpo, base / (tp + tbest + tpo), ts1 / tbest); + { + v2 = true; + std::vector dc2; + chunk_d(0, T, stX, o, dc2); + const double t2 = timeit(dc2), t2a = timeit({dc2[0]}), t2b = timeit({dc2[1]}); + v2 = false; + std::printf(" chunked v2 %.3f (cprep2 %.3f cstate2 %.3f); stage pre+v2+post %.3f = x%.3f vs scanr_4 stage (scan x%.3f)\n", t2, t2a, t2b, + tp + t2 + tpo, base / (tp + t2 + tpo), ts1 / t2); + } + if (std::getenv("GDN_PROBE_RUN")) { + v2 = true; + std::vector dd; + chunk_d(0, T, stX, o, dd); + for (int pr = 1; pr <= 3; ++pr) { + mtl::Pipeline a1 = cx.pipeline(lib, "gdn_cprep2", mtl::Constants().i(6, pr)), a2 = cx.pipeline(lib, "gdn_cstate2", mtl::Constants().i(6, pr)); + mtl::Dispatch d1 = dd[0], d2 = dd[1]; + d1.p = &a1; d2.p = &a2; + const double x1 = timeit({d1}), x2 = timeit({d2}); + std::printf(" probe %d: cprep2 %.3f cstate2 %.3f ms\n", pr, x1, x2); + } + v2 = false; + } + { + double tf[3]; + for (int w = 0; w < 3; ++w) { std::vector d1; fused_d(w, 0, T, stX, d1); tf[w] = timeit(d1); } + const double fb = tp + std::min(tf[0], std::min(tf[1], tf[2])); + std::printf(" fused scan+post TB4 %.3f TB8 %.3f TB2 %.3f ms; stage pre+fused best %.3f = x%.3f vs scanr_4 stage\n", tf[0], tf[1], tf[2], fb, base / fb); + } + std::fflush(stdout); + } + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/gdn_layer_bench.cpp b/edge0-strata/metal/tests/gdn_layer_bench.cpp new file mode 100644 index 00000000..c8893236 --- /dev/null +++ b/edge0-strata/metal/tests/gdn_layer_bench.cpp @@ -0,0 +1,194 @@ +// metal/tests/gdn_layer_bench.cpp - one Flash-Next GDN layer's prompt cost, piece by piece (30 Oct GDN-M phase 2). +// Real shapes (N 2560, qkv 10240, z 6144, alpha/beta 48, out 2560 x 6144), a chunk of T rows (the engine's prompt chunk is +// 2048): dq (Q6_K -> fp16 weights), the four input GEMMs separate and fused into one, tile choices, gdn_pre, the scan, gdn_post, +// the out GEMM; each alone (GPU time) and the layer in order, plus the TFLOP/s of every GEMM. Random weights: a timing +// of the byte and FLOP stream, not a numerical check. +// gdn_layer_bench [--T 2048] [--reps 7] [--scan gdn_scanr_4] +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int N = 2560, HK = 16, HV = 48, CH = 10240, ZD = 6144, S = 128, TQ6K = 14; +struct GemmArgs { int M, N, K, pad; }; +struct DqArgs { uint32_t n_in, n_out, row_bytes, perm; }; +struct GdnPArgs { uint32_t n_tok, h_k, h_v, channels; float eps; uint32_t f32; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +} // namespace + +int main(int argc, char** argv) { + int T = 2048, reps = 7; + std::string scan_name = "gdn_scanr_4"; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--T") T = std::atoi(argv[++i]); + else if (a == "--reps") reps = std::atoi(argv[++i]); + else if (a == "--scan") scan_name = argv[++i]; + } + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal"}); + const int lgdn = cx.compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("device %s, chunk T %d, scan %s\n", cx.device_name().c_str(), T, scan_name.c_str()); + ref::Rng r(5); + auto fill_half = [&](mtl::Buffer& b, size_t n, float sc) { + for (size_t i = 0; i < n; ++i) b.as()[i] = ref::f32_to_f16(r.sym() * sc); + }; + const int rows[5] = {CH, ZD, 48, 48, N}; // qkv, z, alpha, beta, out + const int kin[5] = {N, N, N, N, ZD}; + mtl::Buffer mixh = cx.alloc(2 * (size_t)T * N), ygh = cx.alloc(2 * (size_t)T * ZD); + fill_half(mixh, (size_t)T * N, 1.0f); + fill_half(ygh, (size_t)T * ZD, 1.0f); + mtl::Buffer Wq[5], Wh[5], Y[5]; + mtl::Buffer qsrc[5]; + for (int i = 0; i < 5; ++i) { + Wh[i] = cx.alloc(2 * (size_t)rows[i] * kin[i]); + fill_half(Wh[i], (size_t)rows[i] * kin[i], 0.03f); + Y[i] = cx.alloc(4 * (size_t)T * rows[i]); + const size_t rb = ref::row_bytes(TQ6K, kin[i]); + qsrc[i] = cx.alloc(rb * rows[i]); + for (size_t j = 0; j < rb * rows[i]; ++j) qsrc[i].as()[j] = (uint8_t)(r.uni() * 255.0f); + // keep the fp16 scale of each Q6_K block sane (the last two bytes of a 210 byte block) + for (size_t j = 0; j + 210 <= rb * rows[i]; j += 210) { qsrc[i].as()[(j + 208) / 2] = ref::f32_to_f16(0.01f); } + } + const int NF = CH + ZD + 48 + 48; + mtl::Buffer WF = cx.alloc(2 * (size_t)NF * N), YF = cx.alloc(4 * (size_t)T * NF); + fill_half(WF, (size_t)NF * N, 0.03f); + + auto timeit = [&](const std::vector& ds) { + cx.warm(ds, 0.3); + std::vector t; + for (int i = 0; i < reps; ++i) t.push_back(cx.run(ds)); + return median(t) * 1e3; + }; + auto dq_d = [&](int i) { + DqArgs a{(uint32_t)kin[i], (uint32_t)rows[i], (uint32_t)ref::row_bytes(TQ6K, kin[i]), 0}; + static mtl::Pipeline p = cx.pipeline(lib, "dq_f16_14"); + mtl::Dispatch d; d.p = &p; + d.args = {mtl::Arg::B(qsrc[i]), mtl::Arg::B(Wh[i]), mtl::Arg::V(a)}; + d.grid[0] = (kin[i] / 32 + 63) / 64; d.grid[1] = rows[i]; d.tg[0] = 64; + return d; + }; + auto gemm_d = [&](const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Yb, int M, int Nn, int K, int tm, int tn, int sgn = 4) { + static std::vector> cache; + const std::string fn = "gemm_f16_" + std::to_string(tm) + "_" + std::to_string(tn) + "_" + std::to_string(sgn); + mtl::Pipeline* p = nullptr; + for (auto& c : cache) if (c.first == fn) p = c.second; + if (!p) { p = new mtl::Pipeline(cx.pipeline(lib, fn)); cache.push_back({fn, p}); } + GemmArgs g{M, Nn, K, 0}; + mtl::Dispatch d; d.p = p; + d.args = {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Yb), mtl::Arg::V(g)}; + d.grid[0] = (Nn + tn - 1) / tn; d.grid[1] = (M + tm - 1) / tm; d.tg[0] = 32 * sgn; + return d; + }; + auto tfl = [&](int Nn, int K, double ms) { return 2.0 * T * Nn * K / (ms * 1e-3) / 1e12; }; + + std::printf("== dq (Q6_K -> fp16), ms: qkv %.3f z %.3f alpha %.3f beta %.3f out %.3f\n", timeit({dq_d(0)}), timeit({dq_d(1)}), + timeit({dq_d(2)}), timeit({dq_d(3)}), timeit({dq_d(4)})); + std::printf("== GEMMs (tile TMxTN, 4 simdgroups): ms and TFLOP/s\n"); + struct Sh { const char* nm; int i; }; + struct Tl { int tm, tn, sg; }; + const std::vector tiles = {{64, 32, 4}, {64, 64, 4}, {128, 64, 4}, {64, 128, 4}, {128, 128, 4}, {128, 128, 8}, {128, 64, 8}, + {64, 64, 8}, {256, 64, 8}, {64, 256, 8}, {128, 32, 4}, {32, 128, 4}}; + for (Sh s : {Sh{"qkv", 0}, Sh{"z", 1}, Sh{"alpha", 2}, Sh{"out", 4}}) { + std::printf("%-6s N %5d K %4d:", s.nm, rows[s.i], kin[s.i]); + const mtl::Buffer& A = s.i == 4 ? ygh : mixh; + for (const Tl& tl : tiles) { + if (rows[s.i] <= 64 && tl.tn > 32) continue; + const double ms = timeit({gemm_d(A, Wh[s.i], Y[s.i], T, rows[s.i], kin[s.i], tl.tm, tl.tn, tl.sg)}); + std::printf(" %dx%d/%d %.3f (%.0f)", tl.tm, tl.tn, tl.sg, ms, tfl(rows[s.i], kin[s.i], ms)); + } + std::printf("\n"); + } + std::printf("fused qkv+z+alpha+beta N %d:", NF); + double best_f = 1e9; + for (const Tl& tl : tiles) { + const double ms = timeit({gemm_d(mixh, WF, YF, T, NF, N, tl.tm, tl.tn, tl.sg)}); + best_f = std::min(best_f, ms); + std::printf(" %dx%d/%d %.3f (%.0f)", tl.tm, tl.tn, tl.sg, ms, tfl(NF, N, ms)); + } + std::printf("\n"); + // the four input GEMMs as the engine runs them, in order + const double sep = timeit({gemm_d(mixh, Wh[0], Y[0], T, CH, N, 64, 64), gemm_d(mixh, Wh[1], Y[1], T, ZD, N, 64, 64), + gemm_d(mixh, Wh[2], Y[2], T, 48, N, 64, 32), gemm_d(mixh, Wh[3], Y[3], T, 48, N, 64, 32)}); + std::printf("== input projections: separate (engine tiles) %.3f ms, fused best %.3f ms\n", sep, best_f); + + // pre / scan / post + mtl::Buffer qkv = Y[0], z = Y[1], al = Y[2], be = Y[3], cw = cx.alloc(4 * CH * 4), dt = cx.alloc(4 * HV), sa = cx.alloc(4 * HV), + gam = cx.alloc(4 * S), st = cx.alloc(4 * (size_t)S * HV * S), h0 = cx.alloc(4 * CH * 3), h1 = cx.alloc(4 * CH * 3), + qn = cx.alloc(4 * (size_t)T * HK * S), kn = cx.alloc(4 * (size_t)T * HK * S), vc = cx.alloc(4 * (size_t)T * HV * S), + gb = cx.alloc(8 * (size_t)T * HV), kq = cx.alloc(4 * (size_t)T * HK), o = cx.alloc(4 * (size_t)T * HV * S), + yh = cx.alloc(2 * (size_t)T * HV * S), y32 = cx.alloc(4); + for (int i = 0; i < CH * 4; ++i) cw.as()[i] = r.sym() * 0.6f; + for (int i = 0; i < HV; ++i) { dt.as()[i] = r.sym(); sa.as()[i] = -0.2f - r.uni() * 2.0f; } + for (int i = 0; i < S; ++i) gam.as()[i] = 0.5f + r.uni(); + for (size_t i = 0; i < (size_t)T * CH; ++i) qkv.as()[i] = r.sym() * 2.0f; + for (size_t i = 0; i < (size_t)T * HV; ++i) { al.as()[i] = r.sym() * 3.0f; be.as()[i] = r.sym() * 2.0f; } + for (size_t i = 0; i < (size_t)T * ZD; ++i) z.as()[i] = r.sym() * 2.0f; + mtl::Pipeline pre = cx.pipeline(lgdn, "gdn_pre"), post = cx.pipeline(lgdn, "gdn_post"), + scan = cx.pipeline(lgdn, scan_name, mtl::Constants().i(6, 0)); + GdnPArgs ga{(uint32_t)T, HK, HV, CH, 1e-6f, 0}; + mtl::Dispatch dpre; dpre.p = ⪯ + dpre.args = {mtl::Arg::B(qkv), mtl::Arg::B(h0), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), mtl::Arg::B(sa), + mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), mtl::Arg::B(h1), mtl::Arg::V(ga)}; + dpre.grid[0] = HK + HV / 2; dpre.grid[1] = T; dpre.tg[0] = 256; + mtl::Dispatch dscan; dscan.p = &scan; + dscan.args = {mtl::Arg::B(st), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), mtl::Arg::B(o), + mtl::Arg::V(ga)}; + dscan.grid[0] = HV * 4; dscan.tg[0] = 128; + mtl::Dispatch dpost; dpost.p = &post; + dpost.args = {mtl::Arg::B(o), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh), mtl::Arg::B(y32), mtl::Arg::V(ga)}; + dpost.grid[0] = HV; dpost.grid[1] = T; dpost.tg[0] = 128; + const double tpre = timeit({dpre}), tscan = timeit({dscan}), tpost = timeit({dpost}); + std::printf("== gdn_pre %.3f, scan %.3f, gdn_post %.3f ms (stage %.3f)\n", tpre, tscan, tpost, tpre + tscan + tpost); + double tF = 0; + { // the fused-stride kernels (function constant 7) against the separate ones: bit for bit + for (int t = 0; t < T; ++t) { + std::memcpy(YF.as() + (size_t)t * NF, qkv.as() + (size_t)t * CH, 4 * CH); + std::memcpy(YF.as() + (size_t)t * NF + CH, z.as() + (size_t)t * ZD, 4 * ZD); + std::memcpy(YF.as() + (size_t)t * NF + CH + ZD, al.as() + (size_t)t * 48, 4 * 48); + std::memcpy(YF.as() + (size_t)t * NF + CH + ZD + 48, be.as() + (size_t)t * 48, 4 * 48); + } + std::memset(h1.ptr, 0, 4 * CH * 3); + cx.run({dpre, dscan, dpost}); + std::vector rqn((char*)qn.ptr, (char*)qn.ptr + 4 * (size_t)T * HK * S), rvc((char*)vc.ptr, (char*)vc.ptr + 4 * (size_t)T * HV * S), + rgb((char*)gb.ptr, (char*)gb.ptr + 8 * (size_t)T * HV), ryh((char*)yh.ptr, (char*)yh.ptr + 2 * (size_t)T * HV * S), + rh((char*)h1.ptr, (char*)h1.ptr + 4 * CH * 3); + mtl::Pipeline preF = cx.pipeline(lgdn, "gdn_pre", mtl::Constants().i(7, NF)), postF = cx.pipeline(lgdn, "gdn_post", mtl::Constants().i(7, NF)); + mtl::Dispatch fpre = dpre, fpost = dpost; + fpre.p = &preF; fpost.p = &postF; + fpre.args[0] = mtl::Arg::B(YF, 0); + fpre.args[3] = mtl::Arg::B(YF, 4 * (size_t)(CH + ZD)); + fpre.args[4] = mtl::Arg::B(YF, 4 * (size_t)(CH + ZD + 48)); + fpost.args[1] = mtl::Arg::B(YF, 4 * (size_t)CH); + std::memset(qn.ptr, 0, rqn.size()); std::memset(vc.ptr, 0, rvc.size()); std::memset(gb.ptr, 0, rgb.size()); + std::memset(yh.ptr, 0, ryh.size()); std::memset(h1.ptr, 0, rh.size()); + cx.run({fpre, dscan, fpost}); + std::printf("fused-stride gdn_pre/gdn_post vs separate: qn %s, vc %s, gate/beta %s, history %s, fp16 y %s\n", + std::memcmp(qn.ptr, rqn.data(), rqn.size()) == 0 ? "equal" : "DIFFER", std::memcmp(vc.ptr, rvc.data(), rvc.size()) == 0 ? "equal" : "DIFFER", + std::memcmp(gb.ptr, rgb.data(), rgb.size()) == 0 ? "equal" : "DIFFER", std::memcmp(h1.ptr, rh.data(), rh.size()) == 0 ? "equal" : "DIFFER", + std::memcmp(yh.ptr, ryh.data(), ryh.size()) == 0 ? "equal" : "DIFFER"); + tF = timeit({fpre}) + timeit({fpost}); + std::printf("fused-stride pre + post %.3f ms (separate %.3f)\n", tF, tpre + tpost); + } + const double tout = timeit({gemm_d(ygh, Wh[4], Y[4], T, N, ZD, 64, 64)}); + double dqall = 0; + for (int i = 0; i < 5; ++i) dqall += timeit({dq_d(i)}); + std::printf("== layer in engine order (GPU ms): dq %.3f + input GEMMs %.3f + pre/scan/post %.3f + out GEMM %.3f = %.3f (sum of parts)\n", + dqall, sep, tpre + tscan + tpost, tout, dqall + sep + tpre + tscan + tpost + tout); + { // the layer as the engine now runs it on the M5 family: fused input GEMM, fused-stride pre/post, tiled out GEMM + double tout_best = 1e9; + for (auto [tm, tn] : std::vector>{{64, 64}, {128, 64}, {64, 128}}) + tout_best = std::min(tout_best, timeit({gemm_d(ygh, Wh[4], Y[4], T, N, ZD, tm, tn)})); + const double base_l = dqall + sep + tpre + tscan + tpost + tout; + const double new_l = dqall + best_f + tF + tscan + tout_best; + std::printf("== layer: before %.3f ms, after (fused input GEMM %.3f, fused-stride pre+post %.3f, scan %.3f, out GEMM %.3f) %.3f ms = x%.3f\n", + base_l, best_f, tF, tscan, tout_best, new_l, base_l / new_l); + } + return 0; +} diff --git a/edge0-strata/metal/tests/gdn_parity.cpp b/edge0-strata/metal/tests/gdn_parity.cpp new file mode 100644 index 00000000..7288eb1d --- /dev/null +++ b/edge0-strata/metal/tests/gdn_parity.cpp @@ -0,0 +1,213 @@ +// metal/tests/gdn_parity.cpp - the fused GDN layer step (gdn.metal) against a double-precision CPU reference of the +// native contract, the window properties, and its speed. +// +// gdn_parity parity for windows of 1..4 tokens: y, the committed state and conv history; a verify +// window leaves the state untouched; its first m outputs are bitwise a fold of m +// gdn_parity --bench GPU time per layer: 36 distinct layer states rotated (113 MB, beyond the SLC) +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int S = 128, HK = 16, HV = 48, C = (2 * HK + HV) * S, MAXT = 8; +struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +struct Err { + double se = 0, sr = 0; + void add(double g, double r) { se += (g - r) * (g - r); sr += r * r; } + double l2() const { return sr > 0 ? std::sqrt(se / sr) : std::sqrt(se); } +}; +} // namespace + +int main(int argc, char** argv) { + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + mtl::Context cx; + const int lib = cx.compile({"gdn.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("device %s\n", cx.device_name().c_str()); + const int LAYERS = bench ? 36 : 1; + const size_t st_elems = (size_t)S * HV * S; + mtl::Buffer state = cx.alloc(4 * st_elems * LAYERS), st0 = cx.alloc(4 * st_elems), hist = cx.alloc(4 * C * 3 * LAYERS), + h0 = cx.alloc(4 * C * 3), qkv = cx.alloc(4 * C * MAXT), al = cx.alloc(4 * HV * MAXT), + be = cx.alloc(4 * HV * MAXT), z = cx.alloc(4 * HV * S * MAXT), cw = cx.alloc(4 * C * 4), + dt = cx.alloc(4 * HV), sa = cx.alloc(4 * HV), gm = cx.alloc(4 * S), y = cx.alloc(4 * HV * S * MAXT), + y2 = cx.alloc(4 * HV * S * MAXT), hist2 = cx.alloc(4 * C * 3 * LAYERS); + ref::Rng r(17); + for (size_t i = 0; i < st_elems; ++i) st0.as()[i] = r.sym() * 0.1f; + for (int i = 0; i < C * 3; ++i) h0.as()[i] = r.sym(); + for (int i = 0; i < C * MAXT; ++i) qkv.as()[i] = r.sym() * 2.0f; + for (int i = 0; i < HV * MAXT; ++i) { al.as()[i] = r.sym() * 2.0f; be.as()[i] = r.sym() * 2.0f; } + for (int i = 0; i < HV * S * MAXT; ++i) z.as()[i] = r.sym() * 2.0f; + for (int i = 0; i < C * 4; ++i) cw.as()[i] = r.sym() * 0.5f; + for (int i = 0; i < HV; ++i) { dt.as()[i] = r.sym(); sa.as()[i] = -std::exp(r.sym()); } + for (int i = 0; i < S; ++i) gm.as()[i] = 1.0f + 0.2f * r.sym(); + for (int l = 0; l < LAYERS; ++l) { + std::memcpy(state.as() + l * st_elems, st0.ptr, 4 * st_elems); + std::memcpy(hist.as() + l * C * 3, h0.ptr, 4 * C * 3); + } + auto disp = [&](const mtl::Pipeline& p, int layer, int commit, const mtl::Buffer& out) { + GdnArgs a{HK, HV, C, (uint32_t)commit, 1e-6f, 0}; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(state, 4 * st_elems * layer), mtl::Arg::B(hist, 4 * (size_t)C * 3 * layer), mtl::Arg::B(qkv), + mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(z), mtl::Arg::B(cw), mtl::Arg::B(dt), mtl::Arg::B(sa), + mtl::Arg::B(gm), mtl::Arg::B(out), mtl::Arg::V(a), mtl::Arg::B(hist2, 4 * (size_t)C * 3 * layer)}; + d.grid[0] = HV; + d.tg[0] = std::getenv("GDN_IB8") ? 1024 : 512; + return d; + }; + { + // every specialization the engine launches (GDN_NTOK 1..8): an M2 Max gave 448 for ntok >= 2 (#2323) + const char* kn = std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step"; + const int need = std::getenv("GDN_IB8") ? 1024 : 512; + int low = 0; + for (int nt = 1; nt <= 8; ++nt) { + mtl::Pipeline pn = cx.pipeline(lib, kn, mtl::Constants().i(2, nt)); + std::printf("%s pipeline GDN_NTOK %d: max %d threads per threadgroup (the kernel needs %d)%s\n", kn, nt, pn.max_threads, + need, pn.max_threads < need ? " TOO FEW" : ""); + low += pn.max_threads < need; + } + if (low) { + std::printf("gdn parity: %d specialization(s) cannot run with %d threads on this GPU: FAIL\n", low, need); + return 1; + } + } + if (bench) { + for (int nt : {1, 2, 3, 4}) { + const char* kn = std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step"; + mtl::Pipeline p = cx.pipeline(lib, kn, mtl::Constants().i(2, nt)); + for (int commit : {1, 0}) { + std::vector ds; + for (int k = 0; k < 72; ++k) ds.push_back(disp(p, k % LAYERS, commit, y)); + cx.warm(ds, 0.1); + std::vector t; + for (int i = 0; i < 7; ++i) t.push_back(cx.run(ds) / 72); + const double m = median(t), bytes = 4.0 * st_elems * (commit ? 2 : 1); + std::printf("gdn_step ntok %d %-7s %7.2f us per layer state %.1f MB at %6.1f GB/s (x36 layers: %.3f ms)\n", + nt, commit ? "commit" : "verify", m * 1e6, bytes / 1e6, bytes / m / 1e9, 36 * m * 1e3); + } + } + return 0; + } + int fails = 0; + for (int nt = 1; nt <= 4; ++nt) { + mtl::Pipeline p = cx.pipeline(lib, std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step", mtl::Constants().i(2, nt)); + // determinism: the same verify call three times + { + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); + cx.run({disp(p, 0, 0, y2)}); + std::vector first((float*)y2.ptr, (float*)y2.ptr + HV * S * nt); + int diff = 0; + for (int k = 0; k < 2; ++k) { + cx.run({disp(p, 0, 0, y2)}); + diff += std::memcmp(first.data(), y2.ptr, 4 * HV * S * nt) != 0; + } + if (diff) std::printf(" ntok %d: NOT deterministic (%d of 2 repeats differ)\n", nt, diff); + } + // verify: state untouched + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); + cx.run({disp(p, 0, 0, y2)}); + const bool untouched = std::memcmp(state.ptr, st0.ptr, 4 * st_elems) == 0 && std::memcmp(hist.ptr, h0.ptr, 4 * C * 3) == 0; + cx.run({disp(p, 0, 1, y)}); + const bool same_y = std::memcmp(y.ptr, y2.ptr, 4 * HV * S * nt) == 0; + // the reference + std::vector Sd(st_elems), Hd(C * 3); + for (size_t i = 0; i < st_elems; ++i) Sd[i] = st0.as()[i]; + for (int i = 0; i < C * 3; ++i) Hd[i] = h0.as()[i]; + Err ey, es, eh; + for (int t = 0; t < nt; ++t) { + std::vector cv(C); + for (int c = 0; c < C; ++c) { + const double in = qkv.as()[t * C + c]; + double s = Hd[c * 3] * cw.as()[c * 4] + Hd[c * 3 + 1] * cw.as()[c * 4 + 1] + + Hd[c * 3 + 2] * cw.as()[c * 4 + 2] + in * cw.as()[c * 4 + 3]; + cv[c] = s / (1.0 + std::exp(-s)); + Hd[c * 3] = Hd[c * 3 + 1]; + Hd[c * 3 + 1] = Hd[c * 3 + 2]; + Hd[c * 3 + 2] = in; + } + for (int hh = 0; hh < 2 * HK; ++hh) { + double ss = 0; + for (int i = 0; i < S; ++i) ss += cv[hh * S + i] * cv[hh * S + i]; + const double rs = 1.0 / std::sqrt(ss / S + 1e-6 / S); + for (int i = 0; i < S; ++i) cv[hh * S + i] = cv[hh * S + i] * rs / std::sqrt((double)S); + } + for (int h = 0; h < HV; ++h) { + const int hk = h % HK; + const double* q = &cv[hk * S]; + const double* k = &cv[HK * S + hk * S]; + const double* v = &cv[2 * HK * S + h * S]; + const double av = (double)al.as()[t * HV + h] + dt.as()[h]; + const double sp = av > 20 ? av : std::log1p(std::exp(av)); + const double g = std::exp(sp * sa.as()[h]); + const double b = 1.0 / (1.0 + std::exp(-(double)be.as()[t * HV + h])); + std::vector o(S); + for (int jj = 0; jj < S; ++jj) { + double kv = 0; + for (int i = 0; i < S; ++i) kv += Sd[((size_t)i * HV + h) * S + jj] * k[i]; + const double delta = (v[jj] - g * kv) * b; + double acc = 0; + for (int i = 0; i < S; ++i) { + double& sv = Sd[((size_t)i * HV + h) * S + jj]; + sv = g * sv + k[i] * delta; + acc += sv * q[i]; + } + o[jj] = acc / std::sqrt((double)S); + } + double ss = 0; + for (int jj = 0; jj < S; ++jj) ss += o[jj] * o[jj]; + const double rs = 1.0 / std::sqrt(ss / S + 1e-6); + for (int jj = 0; jj < S; ++jj) { + const double zz = z.as()[t * HV * S + h * S + jj]; + ey.add(y.as()[t * HV * S + h * S + jj], o[jj] * rs * gm.as()[jj] / (1.0 + std::exp(-zz))); + } + } + } + for (size_t i = 0; i < st_elems; ++i) es.add(state.as()[i], Sd[i]); + for (int i = 0; i < C * 3; ++i) eh.add(hist2.as()[i], Hd[i]); + // fold property: a window of nt + 1 tokens (no commit) gives bitwise the same first nt outputs + bool prefix = true; + if (nt < 4) { + mtl::Pipeline p2 = cx.pipeline(lib, std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step", mtl::Constants().i(2, nt + 1)); + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); + cx.run({disp(p2, 0, 0, y2)}); + prefix = std::memcmp(y.ptr, y2.ptr, 4 * HV * S * nt) == 0; + } + // partial commit (commit = 0x100 | k): the state and history after the first k tokens, bitwise those of a + // k-token commit; y bitwise the verify window's + bool partial = true; + for (int k = 1; k < nt; ++k) { + mtl::Pipeline pk = cx.pipeline(lib, std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step", mtl::Constants().i(2, k)); + mtl::Pipeline pn = cx.pipeline(lib, std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step", mtl::Constants().i(2, nt)); + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memcpy(hist.ptr, h0.ptr, 4 * C * 3); + cx.run({disp(pk, 0, 1, y2)}); + std::vector sk(state.as(), state.as() + st_elems), hk(hist2.as(), hist2.as() + C * 3); + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + cx.run({disp(pn, 0, 0, y)}); + std::vector yv(y.as(), y.as() + HV * S * nt); + std::memcpy(state.ptr, st0.ptr, 4 * st_elems); + std::memset(hist2.ptr, 0xFF, 4 * C * 3); + cx.run({disp(pn, 0, 0x100 | k, y2)}); + partial = partial && std::memcmp(state.ptr, sk.data(), 4 * st_elems) == 0 && + std::memcmp(hist2.ptr, hk.data(), 4 * C * 3) == 0 && std::memcmp(y2.ptr, yv.data(), 4 * HV * S * nt) == 0; + } + const bool ok = untouched && same_y && prefix && partial && ey.l2() < 1e-5 && es.l2() < 1e-5 && eh.l2() == 0; + fails += !ok; + std::printf("gdn parity ntok %d: y rel L2 %.2e, state %.2e, history %.1e; verify leaves state: %s, verify y == " + "commit y: %s, first %d of a longer window bitwise: %s, partial commits bitwise: %s %s\n", nt, ey.l2(), es.l2(), + eh.l2(), untouched ? "yes" : "NO", same_y ? "yes" : "NO", nt, prefix ? "yes" : "NO", partial ? "yes" : "NO", + ok ? "PASS" : "FAIL"); + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/gdn_prompt_parity.cpp b/edge0-strata/metal/tests/gdn_prompt_parity.cpp new file mode 100644 index 00000000..24cb85bf --- /dev/null +++ b/edge0-strata/metal/tests/gdn_prompt_parity.cpp @@ -0,0 +1,175 @@ +// metal/tests/gdn_prompt_parity.cpp - the prompt GDN passes (gdn_prompt.metal: gdn_pre -> gdn_scan -> gdn_post) +// against the decode step run over the same T tokens in order (gdn.metal's gdn_step, GDN_NTOK = T, itself checked +// against the CPU reference by gdn_parity): y, the final state and the conv history. The scan reassociates +// o = S_new^T q as g S^T q + delta (k . q), so the match is to rounding, not bitwise. Then the speed for T = 2048. +// gdn_prompt_parity [--bench] +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int HK = 16, HV = 48, CH = 10240, S = 128; +struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; }; +struct GdnPArgs { uint32_t n_tok, h_k, h_v, channels; float eps; uint32_t f32; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +} // namespace + +int main(int argc, char** argv) { + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + mtl::Context cx; + const int lib = cx.compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("device %s\n", cx.device_name().c_str()); + const int TMAX = 2048; + auto F = [&](size_t n) { return cx.alloc(4 * n); }; + mtl::Buffer qkv = F((size_t)TMAX * CH), al = F(TMAX * HV), be = F(TMAX * HV), z = F((size_t)TMAX * HV * S), cw = F(CH * 4), + dt = F(HV), sa = F(HV), gam = F(S), st0 = F(S * HV * S), stA = F(S * HV * S), stB = F(S * HV * S), + h0 = F(CH * 3), hA = F(CH * 3), hA2 = F(CH * 3), hB = F(CH * 3), yA = F((size_t)TMAX * HV * S), + yB = F((size_t)TMAX * HV * S), yh = cx.alloc(2 * (size_t)TMAX * HV * S), qn = F((size_t)TMAX * HK * S), + kn = F((size_t)TMAX * HK * S), vc = F((size_t)TMAX * HV * S), gb = F(2 * TMAX * HV), kq = F(TMAX * HK), + o = F((size_t)TMAX * HV * S); + const size_t NCM = (TMAX + 31) / 32; + mtl::Buffer cU = F((size_t)HV * NCM * 32 * S), cW = F((size_t)HV * NCM * 32 * S), clg = F((size_t)HV * NCM * 32), + cD = F((size_t)HV * NCM * 32 * S); + ref::Rng r(21); + for (size_t i = 0; i < (size_t)TMAX * CH; ++i) qkv.as()[i] = r.sym() * 2.0f; + for (int i = 0; i < TMAX * HV; ++i) { al.as()[i] = r.sym() * 3.0f; be.as()[i] = r.sym() * 2.0f; } + for (size_t i = 0; i < (size_t)TMAX * HV * S; ++i) z.as()[i] = r.sym() * 2.0f; + for (int i = 0; i < CH * 4; ++i) cw.as()[i] = r.sym() * 0.6f; + for (int i = 0; i < HV; ++i) { dt.as()[i] = r.sym(); sa.as()[i] = -0.2f - r.uni() * 2.0f; } + for (int i = 0; i < S; ++i) gam.as()[i] = 0.5f + r.uni(); + for (int i = 0; i < S * HV * S; ++i) st0.as()[i] = r.sym() * 0.1f; + for (int i = 0; i < CH * 3; ++i) h0.as()[i] = r.sym(); + auto prompt_ds = [&](int T, const mtl::Pipeline& pre, const mtl::Pipeline& scan, const mtl::Pipeline& post, int cb, + std::vector& ds) { + GdnPArgs a{(uint32_t)T, HK, HV, CH, 1e-6f, 1u}; + mtl::Dispatch d1; + d1.p = ⪯ + d1.args = {mtl::Arg::B(qkv), mtl::Arg::B(hB), mtl::Arg::B(cw), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(dt), + mtl::Arg::B(sa), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), + mtl::Arg::B(hA2), mtl::Arg::V(a)}; + d1.grid[0] = HK + HV / 2; + d1.grid[1] = T; + d1.tg[0] = 256; + ds.push_back(d1); + if (cb == -100) { // the chunked form: gdn_cprep, gdn_cstate, gdn_cout + const uint32_t nc = (T + 31) / 32; + static mtl::Pipeline pcp = cx.pipeline(lib, "gdn_cprep"), pcs = cx.pipeline(lib, "gdn_cstate"), + pco = cx.pipeline(lib, "gdn_cout"); + mtl::Dispatch e1; e1.p = &pcp; + e1.args = {mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(cU), mtl::Arg::B(cW), + mtl::Arg::B(clg), mtl::Arg::V(a)}; + e1.grid[0] = HV; e1.grid[1] = nc; e1.tg[0] = 256; ds.push_back(e1); + mtl::Dispatch e2; e2.p = &pcs; + e2.args = {mtl::Arg::B(stB), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(cU), mtl::Arg::B(cW), mtl::Arg::B(clg), + mtl::Arg::B(cD), mtl::Arg::B(o), mtl::Arg::V(a)}; + e2.grid[0] = HV * 4; e2.tg[0] = 256; ds.push_back(e2); + mtl::Dispatch e3; e3.p = &pco; + e3.args = {mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(cD), mtl::Arg::B(clg), mtl::Arg::B(o), mtl::Arg::V(a)}; + e3.grid[0] = HV; e3.grid[1] = nc; e3.tg[0] = 128; ds.push_back(e3); + } else { + mtl::Dispatch d2; + d2.p = &scan; + d2.args = {mtl::Arg::B(stB), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(vc), mtl::Arg::B(gb), mtl::Arg::B(kq), + mtl::Arg::B(o), mtl::Arg::V(a)}; + // cb > 0: gdn_scan_ (cb column blocks of 32 a threadgroup); cb < 0: gdn_scanr_<-cb> (-cb row blocks) + d2.grid[0] = cb > 0 ? HV * 4 / cb : HV * 4; + d2.tg[0] = 32 * (cb > 0 ? cb : -cb); + ds.push_back(d2); + } + mtl::Dispatch d3; + d3.p = &post; + d3.args = {mtl::Arg::B(o), mtl::Arg::B(z), mtl::Arg::B(gam), mtl::Arg::B(yh), mtl::Arg::B(yB), mtl::Arg::V(a)}; + d3.grid[0] = HV; + d3.grid[1] = T; + d3.tg[0] = 128; + ds.push_back(d3); + }; + mtl::Pipeline pre = cx.pipeline(lib, "gdn_pre"), post = cx.pipeline(lib, "gdn_post"); + int fails = 0; + if (!bench) + for (int T : {1, 3, 37, 300}) + for (int cb : {1, 4, -2, -4, -8, -9, -100}) { + mtl::Constants pc; + if (cb < 0 && cb != -9) pc.i(6, 0); + mtl::Pipeline scan = cx.pipeline(lib, cb == -100 ? std::string("gdn_scanr_4") : cb == -9 ? std::string("gdn_scan_v") : cb > 0 ? "gdn_scan_" + std::to_string(cb) : "gdn_scanr_" + std::to_string(-cb), cb == -100 ? mtl::Constants().i(6, 0) : pc); + if (cb == -9) cb = -8; + mtl::Pipeline step = cx.pipeline(lib, "gdn_step", mtl::Constants().i(2, T)); + std::copy(st0.as(), st0.as() + S * HV * S, stA.as()); + std::copy(st0.as(), st0.as() + S * HV * S, stB.as()); + std::copy(h0.as(), h0.as() + CH * 3, hA.as()); + std::copy(h0.as(), h0.as() + CH * 3, hB.as()); + GdnArgs da{HK, HV, CH, 1, 1e-6f, 0}; + mtl::Dispatch d; + d.p = &step; + d.args = {mtl::Arg::B(stA), mtl::Arg::B(hA), mtl::Arg::B(qkv), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(z), + mtl::Arg::B(cw), mtl::Arg::B(dt), mtl::Arg::B(sa), mtl::Arg::B(gam), mtl::Arg::B(yA), mtl::Arg::V(da), + mtl::Arg::B(hA2)}; + d.grid[0] = HV; + d.tg[0] = 512; + cx.run({d}); + std::vector hdec(hA2.as(), hA2.as() + CH * 3); + std::vector ds; + prompt_ds(T, pre, scan, post, cb, ds); + cx.run(ds); + auto rel = [](const float* a, const float* b, size_t n) { + double se = 0, sr = 0; + for (size_t i = 0; i < n; ++i) { se += (double)(a[i] - b[i]) * (a[i] - b[i]); sr += (double)b[i] * b[i]; } + return std::sqrt(se / std::max(sr, 1e-30)); + }; + const double ry = rel(yB.as(), yA.as(), (size_t)T * HV * S), + rs = rel(stB.as(), stA.as(), S * HV * S), rh = rel(hA2.as(), hdec.data(), CH * 3); + double ryh = 0, syh = 0; + for (size_t i = 0; i < (size_t)T * HV * S; ++i) { + const double e = ref::f16_to_f32(yh.as()[i]) - yA.as()[i]; + ryh += e * e; + syh += (double)yA.as()[i] * yA.as()[i]; + } + const bool ok = ry < 1e-5 && rs < 1e-5 && rh == 0 && std::sqrt(ryh / syh) < 1e-3; + fails += !ok; + std::printf("T %3d, scan CB %d: y rel %.2e (fp16 %.2e), state rel %.2e, history %s %s\n", T, cb, ry, + std::sqrt(ryh / syh), rs, rh == 0 ? "exact" : "DIFFERS", ok ? "PASS" : "FAIL"); + } + if (fails || !bench) { + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; + } + // speed, T = 2048: the decode step over the chunk vs the three passes (each scan width) + const int T = 2048; + mtl::Pipeline step = cx.pipeline(lib, "gdn_step", mtl::Constants().i(2, T)); + GdnArgs da{HK, HV, CH, 1, 1e-6f, 0}; + mtl::Dispatch d; + d.p = &step; + d.args = {mtl::Arg::B(stA), mtl::Arg::B(hA), mtl::Arg::B(qkv), mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(z), + mtl::Arg::B(cw), mtl::Arg::B(dt), mtl::Arg::B(sa), mtl::Arg::B(gam), mtl::Arg::B(yA), mtl::Arg::V(da), mtl::Arg::B(hA2)}; + d.grid[0] = HV; + d.tg[0] = 512; + auto timeit = [&](const std::vector& ds) { + cx.warm(ds, 0.3); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds)); + return median(t); + }; + std::printf("T %d: decode step over the chunk %.2f ms\n", T, timeit({d}) * 1e3); + for (const char* nm : {"gdn_scanr_4", "chunked", "gdn_scanr_4", "chunked"}) { + const bool ch = std::string(nm) == "chunked"; + const int cb = ch ? -100 : nm[9] == 'v' ? -8 : nm[8] == 'r' ? -(nm[10] - '0') : nm[9] - '0'; + mtl::Constants pc; + if (cb < 0 && nm[9] != 'v') pc.i(6, std::getenv("GDN_PROBE") ? std::atoi(std::getenv("GDN_PROBE")) : 0); + mtl::Pipeline scan = cx.pipeline(lib, ch ? "gdn_scanr_4" : nm, pc); + std::vector ds; + prompt_ds(T, pre, scan, post, cb, ds); + const double all = timeit(ds), tp = timeit({ds[0]}), + ts = timeit(std::vector(ds.begin() + 1, ds.end() - 1)), to = timeit({ds.back()}); + if (ch) std::printf(" chunked parts: cprep %.3f, cstate %.3f, cout %.3f ms\n", timeit({ds[1]}) * 1e3, timeit({ds[2]}) * 1e3, + timeit({ds[3]}) * 1e3); + std::printf("T %d, %-13s: pre %.3f + scan %.3f + post %.3f ms = %.3f ms together\n", T, nm, tp * 1e3, ts * 1e3, + to * 1e3, all * 1e3); + } + return 0; +} diff --git a/edge0-strata/metal/tests/gdn_tprobe.cpp b/edge0-strata/metal/tests/gdn_tprobe.cpp new file mode 100644 index 00000000..80622554 --- /dev/null +++ b/edge0-strata/metal/tests/gdn_tprobe.cpp @@ -0,0 +1,36 @@ +// metal/tests/gdn_tprobe.cpp - times kernels/gdn_tprobe.metal: the five MPP matmul2d products of a GDN chunk step, NC chunk steps a +// threadgroup, 192 threadgroups (h_v * 4). Compare with gdn_scanr_4's stage (chunks of 32 -> T = 32 * NC). +// gdn_tprobe [NC] +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include + +int main(int argc, char** argv) { + const int NC = argc > 1 ? std::atoi(argv[1]) : 64; + mtl::Context cx; + const int lib = cx.compile({"gdn_tprobe.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("device %s, NC %d (T %d)\n", cx.device_name().c_str(), NC, NC * 32); + const size_t NB = (size_t)48 * NC; + mtl::Buffer out = cx.alloc(4 * 192 * 4096), W = cx.alloc(2 * NB * 32 * 128), Q = cx.alloc(2 * NB * 32 * 128), + K = cx.alloc(2 * NB * 32 * 128), M = cx.alloc(2 * NB * 32 * 32); + ref::Rng r(3); + for (auto* b : {&W, &Q, &K}) for (size_t i = 0; i < NB * 32 * 128; ++i) b->as()[i] = ref::f32_to_f16(r.sym() * 0.1f); + for (size_t i = 0; i < NB * 32 * 32; ++i) M.as()[i] = ref::f32_to_f16(r.sym() * 0.1f); + for (int pr : {0, 1}) { + mtl::Pipeline p = cx.pipeline(lib, "gdn_tprobe", pr ? mtl::Constants().i(0, pr) : mtl::Constants()); + mtl::Dispatch d; d.p = &p; + d.args = {mtl::Arg::B(out), mtl::Arg::B(W), mtl::Arg::B(Q), mtl::Arg::B(K), mtl::Arg::B(M), mtl::Arg::V(NC)}; + d.grid[0] = 192; d.tg[0] = 128; + cx.warm({d}, 0.3); + std::vector t; + for (int i = 0; i < 7; ++i) t.push_back(cx.run({d}) * 1e3); + std::sort(t.begin(), t.end()); + std::printf("probe %d (%s): %.3f ms for %d chunk steps x 192 threadgroups (max threads %d, tg mem %d)\n", pr, + pr ? "SIMT parts only" : "with the five products", t[3], NC, p.max_threads, p.static_tg_mem); + } + return 0; +} diff --git a/edge0-strata/metal/tests/gemv_parity.cpp b/edge0-strata/metal/tests/gemv_parity.cpp new file mode 100644 index 00000000..5687527f --- /dev/null +++ b/edge0-strata/metal/tests/gemv_parity.cpp @@ -0,0 +1,333 @@ +// metal/tests/gemv_parity.cpp - the dense GEMV (gemv.metal) per format: parity against double, and speed per tile +// shape on the model's projection shapes. +// +// gemv_parity parity: every format, n_in 2560/640/10240, 1..8 tokens +// gemv_parity --bench [--pack Q2_0|IQ3_XXS] [--ntok N] [--all] speed on the pack's shapes (best tile per shape) +// gemv_parity --one TYPE NIN NOUT [--ntok N] sweep one shape +#include "gguf_lite.h" +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; +struct Cfg { int L, R, S, K; }; +const std::vector kCfgs = {{32, 1, 4, 1}, {32, 2, 4, 1}, {16, 1, 4, 1}, {16, 2, 4, 1}, {8, 1, 4, 1}, + {8, 2, 4, 1}, {4, 2, 4, 1}, {4, 1, 4, 1}, {32, 1, 4, 4}, {32, 2, 4, 2}, + {32, 1, 8, 8}, {16, 1, 8, 4}, {32, 1, 8, 2}, {8, 1, 8, 1}, {1, 1, 4, 1}, + {1, 2, 4, 1}, {2, 1, 4, 1}, {2, 2, 4, 1}, {1, 1, 2, 1}, {4, 4, 4, 1}}; +int rows_per_tg(Cfg c) { return (c.S / c.K) * (32 / c.L) * c.R; } +int unit_of(int t) { + switch (t) { + case ref::Q2_0: case ref::IQ3_S: case ref::IQ3_XXS: case ref::IQ2_XXS: case ref::IQ2_XS: case ref::IQ2_S: + case ref::BF16: case ref::F32: return 16; + default: return 32; + } +} +bool cfg_ok(Cfg c, int type, int n_in) { + const int U = n_in / unit_of(type); + return U % c.K == 0 && U / c.K >= c.L / 2; +} +std::string kname(int t, Cfg c) { + return "gemv_" + std::to_string(t) + "_" + std::to_string(c.L) + "_" + std::to_string(c.R) + "_" + + std::to_string(c.S) + "_" + std::to_string(c.K); +} +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +// which prepared x a format reads (units.metal XK_*) +int xkind(int t) { + switch (t) { + case ref::Q2_0: return 1; + case ref::Q4_K: case ref::Q5_K: return 2; + case ref::Q6_K: case ref::Q3_K: return 3; + default: return 0; + } +} +// x [MAXT][n] -> XP [4][MAXT][n] + S16 [MAXT][n/16] +mtl::Dispatch xprep_dispatch(const mtl::Pipeline& p, const mtl::Buffer& x, const mtl::Buffer& xp, uint32_t n, uint32_t nt) { + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(x), mtl::Arg::B(xp), mtl::Arg::V(n), mtl::Arg::V(nt)}; + d.tg[0] = 256; + d.grid[0] = (n / 16 * nt + 255) / 256; + return d; +} + +int type_of(const std::string& s) { + for (int t : {ref::Q2_0, ref::Q4_0, ref::Q5_0, ref::Q8_0, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::IQ4_NL, + ref::IQ4_XS, ref::IQ3_S, ref::BF16}) + if (s == ref::type_name(t)) return t; + std::fprintf(stderr, "unknown type %s\n", s.c_str()); + std::exit(2); +} + +struct Shape { const char* name; int type, n_in, n_out, per_token; }; // per_token: calls per decoded token + +} // namespace + +int main(int argc, char** argv) { + bool bench = false, all = false; + int ntok = 1; + std::string gguf_path; + std::string pack = "Q2_0"; + std::vector one; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--all") all = true; + else if (a == "--ntok") ntok = std::atoi(argv[++i]); + else if (a == "--pack") pack = argv[++i]; + else if (a == "--gguf") gguf_path = argv[++i]; + else if (a == "--one") { + const int t = type_of(argv[i + 1]); + one.push_back({"one", t, std::atoi(argv[i + 2]), std::atoi(argv[i + 3]), 1}); + i += 3; + bench = true; + all = true; + } else { std::fprintf(stderr, "unknown argument %s\n", a.c_str()); return 2; } + } + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); + const int MAXT = 8; + mtl::Pipeline pprep = cx.pipeline(lib, "xprep"); + + if (!gguf_path.empty()) { + // real tensors: the first type-distinct ones of the file (experts: expert 0's matrix), 1 and 3 tokens + gguf::File f; + std::string err; + if (!f.open(gguf_path, err)) { std::fprintf(stderr, "%s\n", err.c_str()); return 2; } + mtl::Buffer Wf = cx.wrap(f.map, f.map_bytes); + std::map pick; + for (const auto& kv : f.tensors) { + const gguf::Tensor& t = kv.second; + if (t.shape.size() < 2 || t.name.find("token_embd") != std::string::npos || + t.name.find("per_layer") != std::string::npos) + continue; + const int ty = t.type; + const bool gemv_type = ty == ref::BF16 || ty == ref::Q8_0 || ty == ref::Q4_0 || ty == ref::Q5_0 || + ty == ref::IQ4_NL || ty == ref::IQ4_XS || ty == ref::Q2_0 || ty == ref::IQ3_S || + ty == ref::Q3_K || ty == ref::Q4_K || ty == ref::Q5_K || ty == ref::Q6_K; + if (gemv_type && !pick.count(ty)) pick[ty] = &t; + } + int fails = 0; + for (const auto& kv : pick) { + const gguf::Tensor& t = *kv.second; + const int ty = t.type, n_in = (int)t.shape[0], n_out = std::min((int)t.shape[1], 4096); + const size_t rb = ref::row_bytes(ty, n_in); + mtl::Buffer x = cx.alloc(4 * n_in * MAXT), y = cx.alloc(4 * (size_t)n_out * MAXT), + xp = cx.alloc(4 * (4 * n_in * MAXT + n_in / 16 * MAXT)); + ref::Rng r(3); + for (int i = 0; i < n_in * MAXT; ++i) x.as()[i] = r.sym(); + cx.run({xprep_dispatch(pprep, x, xp, n_in, MAXT)}); + const size_t plane = (size_t)4 * n_in * MAXT; + std::vector wf((size_t)n_in * n_out); + ref::dequant(ty, (const uint8_t*)f.map + t.offset, wf.data(), wf.size()); + Cfg c = n_in >= 6144 ? Cfg{32, 1, 8, 2} : n_in <= 640 ? Cfg{8, 1, 4, 1} : Cfg{16, 1, 4, 1}; + if (!cfg_ok(c, ty, n_in)) c = Cfg{32, 1, 4, 1}; + for (int nt : {1, 3}) { + mtl::Pipeline p = cx.pipeline(lib, kname(ty, c), mtl::Constants().i(0, nt)); + GemvArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)nt, (uint32_t)rb, (uint32_t)n_in, (uint32_t)n_out}; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(Wf, t.offset), mtl::Arg::B(xp, plane * xkind(ty)), mtl::Arg::B(xp, plane * 4), + mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = (n_out + rows_per_tg(c) - 1) / rows_per_tg(c); + d.tg[0] = c.S * 32; + cx.run({d}); + double se = 0, sr = 0; + for (int tt = 0; tt < nt; ++tt) + for (int o = 0; o < n_out; ++o) { + double ref = 0; + for (int k = 0; k < n_in; ++k) ref += (double)wf[(size_t)o * n_in + k] * x.as()[tt * n_in + k]; + const double e = y.as()[tt * n_out + o] - ref; + se += e * e; + sr += ref * ref; + } + const double rel = std::sqrt(se / sr); + const bool ok = rel < 1e-5 && std::isfinite(rel); + fails += !ok; + std::printf("real %-7s %-34s %5d x %5d ntok %d: rel L2 %.2e %s\n", ref::type_name(ty), t.name.c_str(), + n_out, n_in, nt, rel, ok ? "PASS" : "FAIL"); + } + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; + } + if (!bench) { + int fails = 0; + for (int t : {ref::F32, ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::IQ4_NL, ref::IQ4_XS, ref::Q2_0, ref::IQ3_S, + ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K}) { + for (int n_in : {2560, 640, 10240}) { + if (n_in % ref::block_elems(t)) continue; + const int n_out = 100; // not a multiple of the rows per threadgroup on purpose + const size_t rb = ref::row_bytes(t, n_in); + mtl::Buffer W = cx.alloc(rb * n_out), x = cx.alloc(4 * n_in * MAXT), y = cx.alloc(4 * n_out * MAXT); + mtl::Buffer xp = cx.alloc(4 * (4 * n_in * MAXT + n_in / 16 * MAXT)); + ref::random_blocks(t, W.as(), n_out, n_in, 31 + t + n_in); + ref::Rng r(7 + n_in); + for (int i = 0; i < n_in * MAXT; ++i) x.as()[i] = r.sym(); + cx.run({xprep_dispatch(pprep, x, xp, n_in, MAXT)}); + const size_t plane = (size_t)4 * n_in * MAXT; + std::vector wf((size_t)n_in * n_out); + ref::dequant(t, W.ptr, wf.data(), wf.size()); + for (int nt : {1, 3, 8}) { + double worst = 0; + int not_exact = 0; + std::string worst_cfg; + for (Cfg c : kCfgs) { + if (!cfg_ok(c, t, n_in)) continue; + mtl::Pipeline p = cx.pipeline(lib, kname(t, c), mtl::Constants().i(0, nt)); + GemvArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)nt, (uint32_t)rb, (uint32_t)n_in, + (uint32_t)n_out}; + std::memset(y.ptr, 0, y.bytes); + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(W), mtl::Arg::B(xp, plane * xkind(t)), mtl::Arg::B(xp, plane * 4), + mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = (n_out + rows_per_tg(c) - 1) / rows_per_tg(c); + d.tg[0] = c.S * 32; + cx.run({d}); + double se = 0, sr = 0; + for (int tt = 0; tt < nt; ++tt) + for (int o = 0; o < n_out; ++o) { + double ref = 0; + for (int k = 0; k < n_in; ++k) ref += (double)wf[(size_t)o * n_in + k] * x.as()[tt * n_in + k]; + const double e = y.as()[tt * n_out + o] - ref; + se += e * e; + sr += ref * ref; + } + const double rel = std::sqrt(se / sr); + if (rel > worst) { worst = rel; worst_cfg = kname(t, c); } + if (nt > 1) { + // every token of a window must be bitwise what a one-token call gives + std::vector ywin((float*)y.ptr, (float*)y.ptr + nt * n_out); + mtl::Pipeline p1 = cx.pipeline(lib, kname(t, c), mtl::Constants().i(0, 1)); + GemvArgs a1 = a; + a1.n_tok = 1; + for (int tt = 0; tt < nt; ++tt) { + mtl::Dispatch d1 = d; + d1.p = &p1; + d1.args = {mtl::Arg::B(W), mtl::Arg::B(xp, plane * xkind(t) + 4 * (size_t)tt * n_in), + mtl::Arg::B(xp, plane * 4 + 4 * (size_t)tt * (n_in / 16)), mtl::Arg::B(y), + mtl::Arg::V(a1)}; + cx.run({d1}); + if (std::memcmp(y.ptr, ywin.data() + tt * n_out, 4 * n_out) != 0) ++not_exact; + } + } + } + const bool ok = worst < 1e-5 && not_exact == 0; + fails += !ok; + std::printf("parity %-7s n_in %5d ntok %d: worst rel L2 over %zu tiles %.2e (%s), window == one " + "token: %s %s\n", + ref::type_name(t), n_in, nt, kCfgs.size(), worst, worst_cfg.c_str(), + nt == 1 ? "-" : not_exact ? "NO" : "bitwise", ok ? "PASS" : "FAIL"); + } + cx.free(W); + cx.free(x); + cx.free(y); + cx.free(xp); + } + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; + } + + // ---- bench + std::vector shapes = one; + if (shapes.empty()) { + if (pack == "Q2_0") + shapes = { + {"hc down (x4/layer)", ref::BF16, 10240, 320, 4 * 48}, {"hc up (x4/layer)", ref::BF16, 320, 10240, 4 * 48}, + {"router", ref::BF16, 2560, 512, 48}, {"gdn qkv", ref::IQ4_XS, 2560, 10240, 13}, + {"gdn qkv", ref::Q3_K, 2560, 10240, 18}, {"gdn qkv", ref::Q4_K, 2560, 10240, 4}, + {"gdn gate", ref::Q3_K, 2560, 6144, 25}, {"gdn out", ref::IQ4_XS, 6144, 2560, 17}, + {"gdn out", ref::Q4_K, 6144, 2560, 11}, {"qsa q", ref::Q3_K, 2560, 12288, 6}, + {"qsa out", ref::Q5_K, 6144, 2560, 4}, {"qsa k", ref::Q3_K, 2560, 512, 3}, + {"shexp gate", ref::Q2_0, 2560, 640, 21}, {"shexp down", ref::Q4_0, 640, 2560, 16}, + {"shexp down", ref::Q2_0, 640, 2560, 16}, {"indexer q", ref::BF16, 2560, 512, 12}, + {"ssm alpha", ref::BF16, 2560, 48, 36}, {"lm head", ref::Q5_K, 2560, 248320, 1}, + {"gdn qkv", ref::Q6_K, 2560, 10240, 1}, + }; + else + shapes = { + {"hc down (x4/layer)", ref::BF16, 10240, 320, 4 * 48}, {"hc up (x4/layer)", ref::BF16, 320, 10240, 4 * 48}, + {"gdn qkv", ref::IQ4_XS, 2560, 10240, 11}, {"gdn qkv", ref::Q6_K, 2560, 10240, 14}, + {"gdn qkv", ref::Q4_K, 2560, 10240, 8}, {"gdn out", ref::Q6_K, 6144, 2560, 19}, + {"gdn gate", ref::IQ3_S, 2560, 6144, 10}, {"gdn gate", ref::IQ4_XS, 2560, 6144, 11}, + {"qsa out", ref::Q6_K, 6144, 2560, 10}, {"qsa q", ref::IQ4_XS, 2560, 12288, 4}, + {"shexp down", ref::IQ4_NL, 640, 2560, 40}, {"shexp up", ref::IQ3_S, 2560, 640, 27}, + {"lm head", ref::Q5_K, 2560, 248320, 1}, + }; + } + std::vector ntoks = {ntok}; + std::printf("%-20s %-7s %6s %7s %5s %9s %9s %8s %s\n", "shape", "type", "n_in", "n_out", "ntok", "us", "GB/s", + "MB", "best tile (L,R,S,K)"); + for (const Shape& s : shapes) { + const size_t rb = ref::row_bytes(s.type, s.n_in), mb = rb * s.n_out; + const int copies = (int)std::max(4, std::min(256, (1ull << 30) / mb)); + mtl::Buffer W = cx.alloc(mb * copies), x = cx.alloc(4 * s.n_in * MAXT), y = cx.alloc(4 * (size_t)s.n_out * MAXT); + { + std::vector tmp(mb); + ref::random_blocks(s.type, tmp.data(), s.n_out, s.n_in, 99); + for (int c = 0; c < copies; ++c) std::memcpy(W.as() + (size_t)c * mb, tmp.data(), mb); + } + ref::Rng r(5); + for (int i = 0; i < s.n_in * MAXT; ++i) x.as()[i] = r.sym(); + mtl::Buffer xp = cx.alloc(4 * (4 * s.n_in * MAXT + s.n_in / 16 * MAXT)); + cx.run({xprep_dispatch(pprep, x, xp, s.n_in, MAXT)}); + const size_t plane = (size_t)4 * s.n_in * MAXT; + for (int nt : ntoks) { + double best = 1e9; + Cfg bc{}; + const bool gx = std::getenv("GEMV_X") != nullptr; // gemvx_* (x staged in threadgroup memory) + const uint32_t tgm = (uint32_t)(4 * nt * (s.n_in + s.n_in / 16) + 15) / 16 * 16; + for (Cfg c : gx ? std::vector{{32, 1, 4, 1}, {32, 2, 4, 1}, {16, 1, 4, 1}, {16, 2, 4, 1}, {8, 1, 4, 1}, + {8, 2, 4, 1}, {4, 2, 4, 1}, {16, 2, 8, 1}, {8, 2, 8, 1}, {16, 4, 4, 1}, + {32, 2, 8, 1}, {2, 2, 4, 1}} + : kCfgs) { + if (!cfg_ok(c, s.type, s.n_in)) continue; + if (!all && !(c.K == 1 || s.n_in >= 6144)) continue; + if (gx && tgm > 32768) continue; + const std::string kn = gx ? "gemvx_" + std::to_string(s.type) + "_" + std::to_string(c.L) + "_" + std::to_string(c.R) + + "_" + std::to_string(c.S) + : kname(s.type, c); + mtl::Pipeline p = cx.pipeline(lib, kn, mtl::Constants().i(0, nt).i(9, std::getenv("GEMV_PROBE") ? std::atoi(std::getenv("GEMV_PROBE")) : 0)); + GemvArgs a{(uint32_t)s.n_in, (uint32_t)s.n_out, (uint32_t)nt, (uint32_t)rb, (uint32_t)s.n_in, + (uint32_t)s.n_out}; + std::vector ds; + const int R = std::max(copies, 16); + for (int k = 0; k < R; ++k) { + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(W, (size_t)(k % copies) * mb), mtl::Arg::B(xp, plane * xkind(s.type)), + mtl::Arg::B(xp, plane * 4), mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = (s.n_out + rows_per_tg(c) - 1) / rows_per_tg(c); + d.tg[0] = c.S * 32; + if (gx) d.tg_mem = tgm; + ds.push_back(d); + } + cx.warm(ds, 0.05); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds) / R); + const double m = median(t); + if (one.size()) std::printf(" tile %2d,%d,%d,%d: %8.2f us %6.1f GB/s\n", c.L, c.R, c.S, c.K, m * 1e6, mb / m / 1e9); + if (m < best) { best = m; bc = c; } + } + std::printf("%-20s %-7s %6d %7d %5d %9.2f %9.1f %8.2f %d,%d,%d,%d\n", s.name, ref::type_name(s.type), s.n_in, + s.n_out, nt, best * 1e6, mb / best / 1e9, mb / 1e6, bc.L, bc.R, bc.S, bc.K); + } + cx.free(W); + cx.free(x); + cx.free(y); + cx.free(xp); + } + return 0; +} diff --git a/edge0-strata/metal/tests/gr2_parity.cpp b/edge0-strata/metal/tests/gr2_parity.cpp new file mode 100644 index 00000000..fdc2337b --- /dev/null +++ b/edge0-strata/metal/tests/gr2_parity.cpp @@ -0,0 +1,246 @@ +// metal/tests/gr2_parity.cpp - the two-dispatch hyper-connection read (gr2.metal: gr_down3 + gr_up3), BF16 and Q8_0 +// weights, against a double-precision reference of fused_gr.hpp's formulas on the SAME weights (the Q8_0 copy +// dequantized for its own check), plus how far the Q8_0 read lands from the BF16 one; and the speed. +// +// gr2_parity parity: 1, 3 and 8 tokens, with and without the folded write, BF16 and Q8_0 +// gr2_parity --bench GPU time of the 2 dispatches (48 weight sets rotated) per tile, vs the 3-dispatch read +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int N = 2560, HC = 4, LR = 320, NIN = HC * N, NR = LR + HC, MAXT = 8; +struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; +struct GrArgs3 { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; uint32_t rb_down, rb_up, comb, topk; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +struct Err { + double se = 0, sr = 0; + void add(double g, double r) { se += (g - r) * (g - r); sr += r * r; } + double l2() const { return sr > 0 ? std::sqrt(se / sr) : std::sqrt(se); } +}; +} // namespace + +int main(int argc, char** argv) { + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "gr.metal", "gr2.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); + const int SETS = bench ? 48 : 1; + const size_t bd = (size_t)LR * NIN * 2, bi = (size_t)HC * NIN * 2, bu = (size_t)NIN * LR * 2; + const size_t qd = (size_t)LR * NIN / 32 * 34, qi = (size_t)HC * NIN / 32 * 34, qu = (size_t)NIN * LR / 32 * 34; + const size_t set = bd + bi + bu + qd + qi + qu + NIN * 4; + mtl::Buffer W = cx.alloc(set * SETS); + ref::Rng rng(3); + { + uint8_t* b = W.as(); + ref::random_blocks(ref::BF16, b, LR, NIN, 11, 0.02f); + ref::random_blocks(ref::BF16, b + bd, HC, NIN, 22, 0.02f); + ref::random_blocks(ref::BF16, b + bd + bi, NIN, LR, 33, 0.1f); + std::vector f((size_t)LR * NIN); + ref::dequant(ref::BF16, b, f.data(), (size_t)LR * NIN); + ref::quantize_q8_0(f.data(), (size_t)LR * NIN, b + bd + bi + bu); + ref::dequant(ref::BF16, b + bd, f.data(), (size_t)HC * NIN); + ref::quantize_q8_0(f.data(), (size_t)HC * NIN, b + bd + bi + bu + qd); + ref::dequant(ref::BF16, b + bd + bi, f.data(), (size_t)NIN * LR); + ref::quantize_q8_0(f.data(), (size_t)NIN * LR, b + bd + bi + bu + qd + qi); + float* wn = (float*)(b + bd + bi + bu + qd + qi + qu); + for (int i = 0; i < NIN; ++i) wn[i] = 1.0f + 0.1f * rng.sym(); + for (int s = 1; s < SETS; ++s) std::memcpy(b + s * set, b, set); + } + mtl::Buffer R0 = cx.alloc(4 * MAXT * NIN), Rout = cx.alloc(4 * MAXT * NIN), bo = cx.alloc(4 * MAXT * N), + inj = cx.alloc(4 * MAXT * HC), Dprev = cx.alloc(4 * MAXT * NR * HC), D = cx.alloc(4 * MAXT * NR * HC), + rs = cx.alloc(4 * MAXT * HC), xp = cx.alloc(4 * (4 * MAXT * N + MAXT * N / 16)), + xr = cx.alloc(4 * (4 * MAXT * N + MAXT * N / 16)), xpb = cx.alloc(4 * 2 * MAXT * N), + R3 = cx.alloc(4 * MAXT * NIN), xn = cx.alloc(4 * MAXT * NIN), lo = cx.alloc(4 * MAXT * LR), + injo = cx.alloc(4 * MAXT * HC); + for (int i = 0; i < MAXT * NIN; ++i) R0.as()[i] = rng.sym() * 2.0f; + for (int i = 0; i < MAXT * N; ++i) bo.as()[i] = rng.sym(); + for (int i = 0; i < MAXT * HC; ++i) inj.as()[i] = rng.sym() * 3.0f; + // the previous read's inject partials: stream 0 carries the whole value + for (int t = 0; t < MAXT; ++t) + for (int c = 0; c < HC; ++c) Dprev.as()[(t * NR + LR + c) * HC] = inj.as()[t * HC + c]; + mtl::Pipeline pprep = cx.pipeline(lib, "xprep"); + int fails = 0; + for (int nt : bench ? std::vector{1, 3} : std::vector{1, 3, 8}) { + for (int qq = 0; qq < 2; ++qq) { + const int q8 = std::getenv("GR_Q8FIRST") ? 1 - qq : qq; + const int F = q8 ? 8 : 30; + for (int tt : bench ? std::vector{1, 2, 3} : std::vector{3}) { + for (int rps : bench ? std::vector{1, 2} : std::vector{1}) { + if (bench && tt > nt) continue; + const char* pre = (!q8 && std::getenv("GR_PRELOAD")) ? "gr_down3b_" : "gr_down3_"; + const std::string dn = !q8 && std::getenv("GR_PRELOAD") + ? std::string(pre) + std::to_string(tt) + "_" + std::to_string(rps) + : std::string(pre) + std::to_string(F) + "_" + std::to_string(tt) + "_" + std::to_string(rps); + mtl::Pipeline pd = cx.pipeline(lib, dn, mtl::Constants().i(1, nt)), + pu = cx.pipeline(lib, "gr_up3_" + std::to_string(F), mtl::Constants().i(1, nt)); + for (int apply = 0; apply < 2; ++apply) { + if (bench && !apply) continue; + GrArgs3 a{N, HC, LR, MAXT, (uint32_t)apply, 1e-6f, (uint32_t)(q8 ? NIN / 32 * 34 : NIN * 2), + (uint32_t)(q8 ? LR / 32 * 34 : LR * 2)}; + auto enc = [&](int s, std::vector& ds) { + const size_t o = (size_t)s * set; + const size_t od = q8 ? o + bd + bi + bu : o, oi = q8 ? o + bd + bi + bu + qd : o + bd, + ou = q8 ? o + bd + bi + bu + qd + qi : o + bd + bi, + ow = o + bd + bi + bu + qd + qi + qu; + mtl::Dispatch d; + d.p = &pd; + d.args = {mtl::Arg::B(W, od), mtl::Arg::B(W, oi), mtl::Arg::B(R0), mtl::Arg::B(Rout), mtl::Arg::B(bo), + mtl::Arg::B(Dprev), mtl::Arg::B(W, ow), mtl::Arg::B(D), mtl::Arg::B(rs), mtl::Arg::V(a)}; + d.grid[0] = (NR + 8 * rps - 1) / (8 * rps); + d.grid[1] = HC; + d.tg[0] = 256; + ds.push_back(d); + mtl::Dispatch u; + u.p = &pu; + u.args = {mtl::Arg::B(W, ou), mtl::Arg::B(D), mtl::Arg::B(Rout), mtl::Arg::B(xp), mtl::Arg::V(a), + mtl::Arg::B(W, ow), mtl::Arg::B(rs)}; + u.grid[0] = N / 16; + u.tg[0] = 256; + ds.push_back(u); + }; + if (bench) { + std::vector ds; + for (int k = 0; k < 96; ++k) enc(k % SETS, ds); + if (const char* only = std::getenv("GR_ONLY")) { + std::vector o; + for (size_t i = (std::string(only) == "up"); i < ds.size(); i += 2) o.push_back(ds[i]); + ds = o; + } + cx.warm(ds, 0.1); + std::vector t; + for (int i = 0; i < 7; ++i) t.push_back(cx.run(ds) / 96); + const double bytes = q8 ? qd + qi + qu : bd + bi + bu; + std::printf("bench ntok %d %-4s TT %d RPS %d: 2 dispatches %7.2f us %6.1f GB/s (%.2f MB)\n", nt, + q8 ? "Q8_0" : "BF16", tt, rps, median(t) * 1e6, bytes / median(t) / 1e9, bytes / 1e6); + continue; + } + std::vector ds; + enc(0, ds); + cx.run(ds); + const uint8_t* b = W.as(); + std::vector wd((size_t)LR * NIN), wi((size_t)HC * NIN), wu((size_t)NIN * LR); + const int ty = q8 ? ref::Q8_0 : ref::BF16; + ref::dequant(ty, b + (q8 ? bd + bi + bu : 0), wd.data(), wd.size()); + ref::dequant(ty, b + (q8 ? bd + bi + bu + qd : bd), wi.data(), wi.size()); + ref::dequant(ty, b + (q8 ? bd + bi + bu + qd + qi : bd + bi), wu.data(), wu.size()); + const float* wn = (const float*)(b + bd + bi + bu + qd + qi + qu); + Err eR, elo, einj, emix; + for (int t = 0; t < nt; ++t) { + std::vector r(NIN), x(NIN); + for (int c = 0; c < HC; ++c) { + const double g = apply ? 2.0 / (1.0 + std::exp(-(double)inj.as()[t * HC + c] / HC)) : 0.0; + double ss = 0; + for (int d = 0; d < N; ++d) { + double v = R0.as()[t * NIN + c * N + d]; + if (apply) v += bo.as()[t * N + d] * g; + r[c * N + d] = v; + ss += v * v; + } + const double rsv = 1.0 / std::sqrt(ss / N + 1e-6); + for (int d = 0; d < N; ++d) x[c * N + d] = r[c * N + d] * wn[c * N + d] * rsv; + } + for (int i = 0; i < NIN; ++i) eR.add(Rout.as()[t * NIN + i], r[i]); + std::vector lg(LR); + for (int k = 0; k < LR; ++k) { + double s = 0; + for (int i = 0; i < NIN; ++i) s += (double)wd[(size_t)k * NIN + i] * x[i]; + const double z = s / HC; + const float* dp = D.as() + (t * NR + k) * HC; + const float zg = (((dp[0] + dp[1]) + dp[2]) + dp[3]) / (float)HC; + lg[k] = zg / (1.0f + std::exp(-zg)); + elo.add(lg[k], z / (1.0 + std::exp(-z))); + } + for (int c = 0; c < HC; ++c) { + double s = 0; + for (int i = 0; i < NIN; ++i) s += (double)wi[(size_t)c * NIN + i] * x[i]; + const float* dp = D.as() + (t * NR + LR + c) * HC; + einj.add(((dp[0] + dp[1]) + dp[2]) + dp[3], s); + } + for (int d = 0; d < N; ++d) { + double m = 0; + for (int c = 0; c < HC; ++c) { + double s = 0; + for (int k = 0; k < LR; ++k) s += (double)wu[(size_t)(c * N + d) * LR + k] * lg[k]; + m += x[c * N + d] / (1.0 + std::exp(-s)); + } + emix.add(xp.as()[t * N + d], m / HC); + } + } + { + mtl::Dispatch d; + d.p = &pprep; + const uint32_t n = N, ntt = MAXT; + d.args = {mtl::Arg::B(xp), mtl::Arg::B(xr), mtl::Arg::V(n), mtl::Arg::V(ntt)}; + d.tg[0] = 256; + d.grid[0] = (N / 16 * MAXT + 255) / 256; + cx.run({d}); + } + bool planes = true; + for (int k = 0; k < 4; ++k) + planes = planes && std::memcmp(xp.as() + (size_t)k * MAXT * N, + xr.as() + (size_t)k * MAXT * N, 4 * (size_t)nt * N) == 0; + planes = planes && std::memcmp(xp.as() + 4 * MAXT * N, xr.as() + 4 * MAXT * N, + 4 * nt * N / 16) == 0; + double qdiff = -1; + if (q8) { + Err e; + for (int i = 0; i < nt * N; ++i) e.add(xp.as()[i], xpb.as()[apply * MAXT * N + i]); + qdiff = e.l2(); + } else { + std::memcpy(xpb.as() + apply * MAXT * N, xp.ptr, 4 * (size_t)MAXT * N); + } + const bool ok = eR.l2() < 1e-6 && elo.l2() < 1e-5 && einj.l2() < 1e-5 && emix.l2() < 1e-5 && planes; + fails += !ok; + char extra[96] = ""; + if (qdiff >= 0) std::snprintf(extra, sizeof extra, ", mixed vs the BF16 read %.2e", qdiff); + std::printf("parity ntok %d %-4s apply %d: R' %.2e, lo %.2e, inject %.2e, mixed %.2e, planes %s%s %s\n", + nt, q8 ? "Q8_0" : "BF16", apply, eR.l2(), elo.l2(), einj.l2(), emix.l2(), + planes ? "bitwise" : "NO", extra, ok ? "PASS" : "FAIL"); + } + } + } + } + } + if (bench) { + mtl::Pipeline pn = cx.pipeline(lib, "gr_norm"), pd = cx.pipeline(lib, "gr_down_8_8", mtl::Constants().i(1, 1)), + pu = cx.pipeline(lib, "gr_up", mtl::Constants().i(1, 1)); + GrArgs a{N, HC, LR, 1, 1, 1e-6f}; + std::vector ds; + for (int k = 0; k < 96; ++k) { + const size_t o = (size_t)(k % SETS) * set; + mtl::Dispatch d1, d2, d3; + d1.p = &pn; + d1.args = {mtl::Arg::B(R3), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(W, o + bd + bi + bu + qd + qi + qu), + mtl::Arg::B(xn), mtl::Arg::V(a)}; + d1.grid[0] = HC; + d1.tg[0] = 256; + d2.p = &pd; + d2.args = {mtl::Arg::B(W, o), mtl::Arg::B(W, o + bd), mtl::Arg::B(xn), mtl::Arg::B(lo), mtl::Arg::B(injo), + mtl::Arg::V(a)}; + d2.grid[0] = LR + HC; + d2.tg[0] = 256; + d3.p = &pu; + d3.args = {mtl::Arg::B(W, o + bd + bi), mtl::Arg::B(lo), mtl::Arg::B(xn), mtl::Arg::B(xp), mtl::Arg::V(a)}; + d3.grid[0] = N / 16; + d3.tg[0] = 256; + ds.push_back(d1); + ds.push_back(d2); + ds.push_back(d3); + } + cx.warm(ds, 0.1); + std::vector t; + for (int i = 0; i < 7; ++i) t.push_back(cx.run(ds) / 96); + std::printf("bench ntok 1 BF16 3-dispatch read (gr.metal): %.2f us\n", median(t) * 1e6); + return 0; + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/gr_parity.cpp b/edge0-strata/metal/tests/gr_parity.cpp new file mode 100644 index 00000000..0df1c528 --- /dev/null +++ b/edge0-strata/metal/tests/gr_parity.cpp @@ -0,0 +1,208 @@ +// metal/tests/gr_parity.cpp - the hyper-connection read (gr.metal) against a double-precision CPU reference +// (fused_gr.hpp's formulas), and its speed. +// +// gr_parity [--ntok N] parity for N = 1, 3, 8 (or N), with and without the folded write +// gr_parity --bench [--ntok N] per-kernel GPU time, 8 weight sets rotated (no cache reuse across calls) +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int N = 2560, HC = 4, LR = 320, NIN = HC * N; +struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +struct Err { + double se = 0, sr = 0, mx = 0, mr = 0; + void add(double g, double r) { se += (g - r) * (g - r); sr += r * r; mx = std::max(mx, std::fabs(g - r)); mr = std::max(mr, std::fabs(r)); } + double l2() const { return sr > 0 ? std::sqrt(se / sr) : std::sqrt(se); } +}; +} // namespace + +int main(int argc, char** argv) { + bool bench = false; + std::vector ntoks = {1, 3, 8}; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--bench") bench = true; + else if (a == "--ntok") ntoks = {std::atoi(argv[++i])}; + } + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "gr.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); + const int SETS = bench ? 48 : 1; + const bool q8 = std::getenv("GR_Q8") != nullptr; // the Q8_0 form of the HC matrices (gr_down_q8_*, gr_up_q8) + const int wty = q8 ? ref::Q8_0 : ref::BF16; + // GR_INJ_F32 (with GR_Q8): the inject rows as F32 (the Unsloth files), read by gr_down_q8f_* + const bool injf = q8 && std::getenv("GR_INJ_F32") != nullptr; + const size_t wd_bytes = (size_t)LR * NIN * 2, wi_bytes = (size_t)HC * NIN * 4, wu_bytes = (size_t)NIN * LR * 2; + const size_t set_bytes = wd_bytes + wi_bytes + wu_bytes + NIN * 4; + mtl::Buffer W = cx.alloc(set_bytes * SETS); + ref::Rng rng(3); + for (int s = 0; s < SETS; ++s) { + uint8_t* b = W.as() + s * set_bytes; + ref::random_blocks(ref::BF16, b, LR, NIN, 11 + s, 0.02f); + ref::random_blocks(ref::BF16, b + wd_bytes, HC, NIN, 22 + s, 0.02f); + ref::random_blocks(ref::BF16, b + wd_bytes + wi_bytes, NIN, LR, 33 + s, 0.1f); + float* wn = (float*)(b + wd_bytes + wi_bytes + wu_bytes); + for (int i = 0; i < NIN; ++i) wn[i] = 1.0f + 0.1f * rng.sym(); + if (q8) { + // GR_Q8: the three matrices as Q8_0 (quantized from the BF16 values), each at its BF16 offset (it fits) + const size_t offs[3] = {0, wd_bytes, wd_bytes + wi_bytes}; + const size_t ns[3] = {(size_t)LR * NIN, (size_t)HC * NIN, (size_t)NIN * LR}; + for (int m = 0; m < 3; ++m) { + std::vector f(ns[m]); + ref::dequant(ref::BF16, b + offs[m], f.data(), ns[m]); + if (m == 1 && injf) std::memcpy(b + offs[m], f.data(), 4 * ns[m]); + else ref::quantize_q8_0(f.data(), ns[m], b + offs[m]); + } + } + } + const int MAXT = 8; + // GR_DOWN=SG_KS selects gr_down's variant (default 8_8: 8 simdgroups split one row) + const std::string grd = std::getenv("GR_DOWN") ? std::getenv("GR_DOWN") : "8_8"; + int grd_sg = 8, grd_ks = 8; + std::sscanf(grd.c_str(), "%d_%d", &grd_sg, &grd_ks); + const int grd_rs = grd_sg / grd_ks; + mtl::Buffer R = cx.alloc(4 * MAXT * NIN), R0 = cx.alloc(4 * MAXT * NIN), bo = cx.alloc(4 * MAXT * N), + inj = cx.alloc(4 * MAXT * HC), xn = cx.alloc(4 * MAXT * NIN), lo = cx.alloc(4 * MAXT * LR), + injo = cx.alloc(4 * MAXT * HC), xp = cx.alloc(4 * (4 * MAXT * N + MAXT * N / 16)), + xr = cx.alloc(4 * (4 * MAXT * N + MAXT * N / 16)); + for (int i = 0; i < MAXT * NIN; ++i) R0.as()[i] = rng.sym() * 2.0f; + for (int i = 0; i < MAXT * N; ++i) bo.as()[i] = rng.sym(); + for (int i = 0; i < MAXT * HC; ++i) inj.as()[i] = rng.sym() * 3.0f; + mtl::Pipeline pprep = cx.pipeline(lib, "xprep"); + int fails = 0; + for (int nt : ntoks) { + mtl::Pipeline pn = cx.pipeline(lib, "gr_norm"), + pd = cx.pipeline(lib, std::string(injf ? "gr_down_q8f_" : q8 ? "gr_down_q8_" : "gr_down_") + grd, mtl::Constants().i(1, nt)), + pu = cx.pipeline(lib, q8 ? "gr_up_q8" : "gr_up", mtl::Constants().i(1, nt)); + for (int apply = 0; apply < 2; ++apply) { + GrArgs a{N, HC, LR, (uint32_t)nt, (uint32_t)apply, 1e-6f}; + std::memcpy(R.ptr, R0.ptr, R.bytes); + auto enc = [&](int set, std::vector& ds, int which) { + const size_t o = (size_t)set * set_bytes; + mtl::Dispatch d; + if (which & 1) { + d.p = &pn; + d.args = {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(W, o + wd_bytes + wi_bytes + wu_bytes), + mtl::Arg::B(xn), mtl::Arg::V(a)}; + d.grid[0] = HC; + d.grid[1] = nt; + d.tg[0] = 256; + ds.push_back(d); + } + if (which & 2) { + mtl::Dispatch e; + e.p = &pd; + e.args = {mtl::Arg::B(W, o), mtl::Arg::B(W, o + wd_bytes), mtl::Arg::B(xn), mtl::Arg::B(lo), + mtl::Arg::B(injo), mtl::Arg::V(a)}; + e.grid[0] = (LR + HC + grd_rs - 1) / grd_rs; // RS = SG / KS rows per threadgroup + e.tg[0] = 32 * grd_sg; + ds.push_back(e); + } + if (which & 4) { + mtl::Dispatch f; + f.p = &pu; + f.args = {mtl::Arg::B(W, o + wd_bytes + wi_bytes), mtl::Arg::B(lo), mtl::Arg::B(xn), mtl::Arg::B(xp), + mtl::Arg::V(a)}; + f.grid[0] = N / 16; + f.tg[0] = 256; + ds.push_back(f); + } + }; + if (!bench) { + std::vector ds; + enc(0, ds, 7); + cx.run(ds); + // reference + const uint8_t* b = W.as(); + std::vector wd((size_t)LR * NIN), wi((size_t)HC * NIN), wu((size_t)NIN * LR); + ref::dequant(wty, b, wd.data(), wd.size()); + ref::dequant(injf ? ref::F32 : wty, b + wd_bytes, wi.data(), wi.size()); + ref::dequant(wty, b + wd_bytes + wi_bytes, wu.data(), wu.size()); + const float* wn = (const float*)(b + wd_bytes + wi_bytes + wu_bytes); + Err eR, elo, einj, emix; + for (int t = 0; t < nt; ++t) { + std::vector r(NIN), x(NIN), l(LR); + for (int c = 0; c < HC; ++c) { + const double g = apply ? 2.0 / (1.0 + std::exp(-(double)inj.as()[t * HC + c] / HC)) : 0.0; + double ss = 0; + for (int d = 0; d < N; ++d) { + double v = R0.as()[t * NIN + c * N + d]; + if (apply) v += bo.as()[t * N + d] * g; + r[c * N + d] = v; + ss += v * v; + } + const double rs = 1.0 / std::sqrt(ss / N + 1e-6); + for (int d = 0; d < N; ++d) x[c * N + d] = r[c * N + d] * wn[c * N + d] * rs; + } + for (int i = 0; i < NIN; ++i) eR.add(R.as()[t * NIN + i], r[i]); + for (int k = 0; k < LR; ++k) { + double s = 0; + for (int i = 0; i < NIN; ++i) s += (double)wd[(size_t)k * NIN + i] * x[i]; + const double z = s / HC; + l[k] = z / (1.0 + std::exp(-z)); + elo.add(lo.as()[t * LR + k], l[k]); + } + for (int c = 0; c < HC; ++c) { + double s = 0; + for (int i = 0; i < NIN; ++i) s += (double)wi[(size_t)c * NIN + i] * x[i]; + einj.add(injo.as()[t * HC + c], s); + } + for (int d = 0; d < N; ++d) { + double m = 0; + for (int c = 0; c < HC; ++c) { + double s = 0; + // the up projection reads the GPU's lo, so the check isolates gr_up + for (int k = 0; k < LR; ++k) s += (double)wu[(size_t)(c * N + d) * LR + k] * lo.as()[t * LR + k]; + m += x[c * N + d] / (1.0 + std::exp(-s)); + } + emix.add(xp.as()[t * N + d], m / HC); + } + } + // the prepared planes must be bitwise xprep's of the raw plane + { + mtl::Dispatch d; + d.p = &pprep; + const uint32_t n = N, ntt = nt; + d.args = {mtl::Arg::B(xp), mtl::Arg::B(xr), mtl::Arg::V(n), mtl::Arg::V(ntt)}; + d.tg[0] = 256; + d.grid[0] = (N / 16 * nt + 255) / 256; + cx.run({d}); + } + const size_t elems = (size_t)4 * nt * N + (size_t)nt * N / 16; + const bool planes = std::memcmp(xp.ptr, xr.ptr, 4 * elems) == 0; + const bool ok = eR.l2() < 1e-6 && elo.l2() < 1e-5 && einj.l2() < 1e-5 && emix.l2() < 1e-5 && planes; + fails += !ok; + std::printf("parity ntok %d apply %d: R rel L2 %.2e, lo %.2e, inject %.2e, mixed %.2e, prepared planes " + "== xprep: %s %s\n", nt, apply, eR.l2(), elo.l2(), einj.l2(), emix.l2(), + planes ? "bitwise" : "NO", ok ? "PASS" : "FAIL"); + } else { + for (int which : {1, 2, 4, 7}) { + std::vector ds; + for (int r = 0; r < 96; ++r) enc(r % SETS, ds, which); + cx.warm(ds, 0.1); + std::vector ts; + for (int i = 0; i < 7; ++i) ts.push_back(cx.run(ds) / 96); + const double m = median(ts); + const double bytes = which == 1 ? 2.0 * nt * NIN * 4 + NIN * 4 + nt * NIN * 4 + : which == 2 ? wd_bytes + wi_bytes + : which == 4 ? wu_bytes + : wd_bytes + wi_bytes + wu_bytes; + std::printf("bench ntok %d apply %d %-8s %7.2f us %6.1f GB/s (%.2f MB)\n", nt, apply, + which == 1 ? "norm" : which == 2 ? "down" : which == 4 ? "up" : "all 3", m * 1e6, + bytes / m / 1e9, bytes / 1e6); + } + } + } + } + if (!bench) std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/gridbar_bench.cpp b/edge0-strata/metal/tests/gridbar_bench.cpp new file mode 100644 index 00000000..6420be4e --- /dev/null +++ b/edge0-strata/metal/tests/gridbar_bench.cpp @@ -0,0 +1,66 @@ +// metal/tests/gridbar_bench.cpp - phases of a streaming read separated by dispatch boundaries vs by an in-kernel +// grid barrier (gridbar.metal), and whether a phase sees the previous phase's writes of the other threadgroups. +// gridbar_bench [MiB per phase] [phases] +#include "mtl_rt.h" + +#include +#include +#include +#include + +int main(int argc, char** argv) { + const double mib = argc > 1 ? std::atof(argv[1]) : 4.5; + const int phases = argc > 2 ? std::atoi(argv[2]) : 96; + mtl::Context cx; + const int lib = cx.compile({"gridbar.metal"}, std::getenv("GB_FENCE") ? "#define NO_GGML_COMMON 1\n#define GB_FENCE 1\n" + : "#define NO_GGML_COMMON 1\n"); + mtl::Pipeline pp = cx.pipeline(lib, "k_phases"), p1 = cx.pipeline(lib, "k_phase1"); + const uint32_t n = (uint32_t)(mib * 1048576 / 16); + mtl::Buffer src = cx.alloc((size_t)n * 16 * phases), out = cx.alloc(4u << 20), bar = cx.alloc(4), + flag = cx.alloc(4 * 64 * (size_t)phases), bad = cx.alloc(4); + for (size_t i = 0; i < src.bytes / 4; i += 1024) src.as()[i] = (uint32_t)i; + std::printf("device %s: %d phases of %.1f MiB\n", cx.device_name().c_str(), phases, mib); + // dispatch boundaries + for (uint32_t groups : {40u, 160u, 640u}) { + std::vector ds; + for (int p = 0; p < phases; ++p) { + mtl::Dispatch d; + d.p = &p1; + const uint32_t pu = p; + d.args = {mtl::Arg::B(src), mtl::Arg::B(out), mtl::Arg::V(n), mtl::Arg::V(pu)}; + d.grid[0] = groups; + d.tg[0] = 1024; + ds.push_back(d); + } + cx.warm(ds, 0.2); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds)); + std::sort(t.begin(), t.end()); + std::printf("dispatch per phase, %4u x 1024 threads: %7.2f us a phase %6.1f GB/s\n", groups, + t[2] / phases * 1e6, n * 16.0 * phases / t[2] / 1e9); + } + // grid barrier, at most one threadgroup a core + for (uint32_t groups : {10u, 20u, 40u}) { + mtl::Dispatch d; + d.p = &pp; + const uint32_t ph = phases; + d.args = {mtl::Arg::B(src), mtl::Arg::B(out), mtl::Arg::B(bar), mtl::Arg::B(flag), mtl::Arg::B(bad), + mtl::Arg::V(n), mtl::Arg::V(ph)}; + d.grid[0] = groups; + d.tg[0] = 1024; + std::vector t; + uint32_t badc = 0; + for (int i = 0; i < 6; ++i) { + bar.as()[0] = 0; + bad.as()[0] = 0; + for (size_t k = 0; k < flag.bytes / 4; ++k) flag.as()[k] = 0xFFFFFFFFu; + const double s = cx.run({d}); + if (i) t.push_back(s); + badc += bad.as()[0]; + } + std::sort(t.begin(), t.end()); + std::printf("grid barrier, %4u x 1024 threads: %7.2f us a phase %6.1f GB/s stale flags seen: %u\n", + groups, t[2] / phases * 1e6, n * 16.0 * phases / t[2] / 1e9, badc); + } + return 0; +} diff --git a/edge0-strata/metal/tests/gyro_parity.sh b/edge0-strata/metal/tests/gyro_parity.sh new file mode 100644 index 00000000..31bcf800 --- /dev/null +++ b/edge0-strata/metal/tests/gyro_parity.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# metal/tests/gyro_parity.sh - the Metal kernels of agentionai's Gyro experts (TQ2_T, TQK6, TQK7 and the Hadamard rotation) +# against the CPU references, on the real tensors of a Gyro GGUF. +# +# metal/tests/gyro_parity.sh BUILD_DIR GGUF PACK_DIR [LAYERS...] [--experts-first N] +# +# BUILD_DIR the metal build (cmake -S metal -B build-metal; ninja expert_parity moe_mm_parity) +# GGUF a Gyro file; a sparse copy works (header + the experts of the layers named, see docs/GYRO.md "Testing a file that +# does not fit": the first 16 experts of each layer with --experts-first 16) +# PACK_DIR the pack tools/iq_pack.py --compat-bf16 made from it (hadamard.txt carries the rotation's signs) +# +# 1. had-test the Metal rotation == the CPU's hadamard_rotate, bit for bit (f32, f32 in place, f16, f16 rows) +# 2. decode-check ref_quant's trellis decoder == the format's definition written out bit by bit (sampled blocks of every tensor) +# 3. parity moe_plan / moe_gu / moe_down / moe_combine with the rotation of x and h, 1, 3 and 8 tokens, against double +# 4. prompt mm_gu2 / mm_down2 (the prompt chunk's grouped GEMMs, fp16 operands) for the file's types +set -u +B=$1; G=$2; P=$3; shift 3 +LAYERS=() +EXTRA=() +while [ $# -gt 0 ]; do + case $1 in + --experts-first) EXTRA+=(--experts-first "$2"); shift 2 ;; + *) LAYERS+=("$1"); shift ;; + esac +done +[ ${#LAYERS[@]} -eq 0 ] && LAYERS=(0 47) +rc=0 +run() { echo "+ $*"; "$@" || rc=1; } +run "$B/expert_parity" --had-test --had-pack "$P" +for L in "${LAYERS[@]}"; do + run "$B/expert_parity" --gguf "$G" --layer "$L" --decode-check + for nt in 1 3 8; do + run "$B/expert_parity" --gguf "$G" --layer "$L" --had-pack "$P" --ntok "$nt" --seed "$nt" ${EXTRA[@]+"${EXTRA[@]}"} + done +done +run "$B/moe_mm_parity" --v2 --gu TQK6 --dn TQK7 --T 300 --tile 16,64,16,2 +run "$B/moe_mm_parity" --v2 --gu TQ2_T --dn TQ2_T --T 300 --tile 16,64,16,2 +[ $rc -eq 0 ] && echo "gyro_parity: all PASS" || echo "gyro_parity: FAILED" +exit $rc diff --git a/edge0-strata/metal/tests/mm_bench.cpp b/edge0-strata/metal/tests/mm_bench.cpp new file mode 100644 index 00000000..51f5e8f7 --- /dev/null +++ b/edge0-strata/metal/tests/mm_bench.cpp @@ -0,0 +1,116 @@ +// metal/tests/mm_bench.cpp - matrix throughput (TFLOP/s) of MPP matmul2d and simdgroup_matrix, fp16/bf16/fp32 +// operands, fp32 results, on a square 4096^3 and on prompt-expert shapes; checks one product against the CPU. +// mm_bench +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +struct MmArgs { int M, N, K; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +} // namespace + +int main() { + mtl::Context cx; + std::printf("device %s, metal4 %d\n", cx.device_name().c_str(), cx.supports_metal4()); + const int lib = cx.compile({"mm_bench.metal"}, "#define NO_GGML_COMMON 1\n"); + const int MAX = 4096; + mtl::Buffer A = cx.alloc((size_t)MAX * MAX * 4), B = cx.alloc((size_t)MAX * MAX * 4), C = cx.alloc((size_t)MAX * MAX * 4); + ref::Rng r(1); + // fill as half and float views of the same random values (each kernel reads its own type) + auto fill = [&](int elem) { + for (size_t i = 0; i < (size_t)MAX * MAX; ++i) { + const float v = r.sym(); + if (elem == 4) { A.as()[i] = v; B.as()[i] = r.sym(); } + else if (elem == 2) { A.as()[i] = ref::f32_to_f16(v); B.as()[i] = ref::f32_to_f16(r.sym()); } + else { A.as()[i] = ref::f32_to_bf16(v); B.as()[i] = ref::f32_to_bf16(r.sym()); } + } + }; + struct K { std::string name; int elem; int tm, tn, sg; bool mpp; }; + const std::vector ks = { + {"mpp_mm_half_64_32_4", 2, 64, 32, 4, true}, {"mpp_mm_half_64_64_4", 2, 64, 64, 4, true}, + {"mpp_mm_half_128_64_4", 2, 128, 64, 4, true}, {"mpp_mm_half_32_64_4", 2, 32, 64, 4, true}, + {"mpp_mm_half_64_128_4", 2, 64, 128, 4, true}, {"mpp_mm_half_128_128_4", 2, 128, 128, 4, true}, + {"mpp_mm_half_16_64_2", 2, 16, 64, 2, true}, {"mpp_mm_bfloat_64_64_4", 3, 64, 64, 4, true}, + {"mpp_mm_bfloat_128_64_4", 3, 128, 64, 4, true}, {"mpp_mm_float_64_32_4", 4, 64, 32, 4, true}, + {"mpp_mm_float_64_64_4", 4, 64, 64, 4, true}, {"mpp_mm_float_128_64_4", 4, 128, 64, 4, true}, + {"sg_mm_half", 2, 64, 64, 4, false}, {"sg_mm_float", 4, 64, 64, 4, false}, + {"mpp_mm_half_64_16_2", 2, 64, 16, 2, true}, {"mpp_mm_half_64_16_4", 2, 64, 16, 4, true}, + {"mpp_mm_half_128_16_4", 2, 128, 16, 4, true}, {"mpp_mm_half_128_32_4", 2, 128, 32, 4, true}, + {"mpp_mm_half_64_32_2", 2, 64, 32, 2, true}, {"mpp_mm_half_32_16_2", 2, 32, 16, 2, true}, + {"mpp_mm_half_32_32_2", 2, 32, 32, 2, true}, {"mpp_mm_half_16_32_1", 2, 16, 32, 1, true}, + {"mpp_mm_half_16_16_1", 2, 16, 16, 1, true}, {"mpp_mm_half_128_48_4", 2, 128, 48, 4, true}, + {"mpp_mm_half_64_48_4", 2, 64, 48, 4, true}}; + struct Shape { const char* name; int M, N, K; }; + const std::vector shapes = {{"square 4096^3", 4096, 4096, 4096}, + {"expert gu 128 tok x 1280 x 2560", 128, 1280, 2560}, + {"expert down 128 tok x 2560 x 640", 128, 2560, 640}, + {"dense qkv 2048 tok x 10240 x 2560", 2048, 4096, 2560}, + {"expert gu 16 tok x 1280 x 2560", 16, 1280, 2560}, + {"expert gu^T 1280 x 48 tok x 2560", 1280, 48, 2560}, + {"expert gu^T 1280 x 16 tok x 2560", 1280, 16, 2560}, + {"many gu^T 12800 x 48 tok x 2560", 12800, 48, 2560}, + {"many gu 48 tok x 12800 x 2560", 48, 12800, 2560}}; + int filled = 0; + const char* only = std::getenv("MMB_ONLY"); // a substring of the kernel names to run + for (const K& k : ks) { + if (only && k.name.find(only) == std::string::npos) continue; + if (filled != k.elem) { fill(k.elem); filled = k.elem; } + mtl::Pipeline p = cx.pipeline(lib, k.name); + // correctness on a small product (M = tm*2, N = tn*2, K = 256) against the CPU + { + const int M = k.tm * 2, N = k.tn * 2, KK = 256; + MmArgs a{M, N, KK}; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(C), mtl::Arg::V(a)}; + d.grid[0] = N / k.tn; + d.grid[1] = M / k.tm; + d.tg[0] = 32 * k.sg; + cx.run({d}); + double se = 0, sr = 0; + for (int i = 0; i < M; ++i) + for (int j = 0; j < N; ++j) { + double ref = 0; + for (int q = 0; q < KK; ++q) { + auto val = [&](const mtl::Buffer& X, size_t idx) -> double { + if (k.elem == 4) return X.as()[idx]; + if (k.elem == 2) return ref::f16_to_f32(X.as()[idx]); + return ref::bf16_to_f32(X.as()[idx]); + }; + ref += val(A, (size_t)i * KK + q) * val(B, (size_t)j * KK + q); + } + const double e = C.as()[i * N + j] - ref; + se += e * e; + sr += ref * ref; + } + std::printf("%-26s check rel L2 %.2e\n", k.name.c_str(), std::sqrt(se / sr)); + } + for (const Shape& s : shapes) { + if (s.M % k.tm || s.N % k.tn) continue; + MmArgs a{s.M, s.N, s.K}; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(C), mtl::Arg::V(a)}; + d.grid[0] = s.N / k.tn; + d.grid[1] = s.M / k.tm; + d.tg[0] = 32 * k.sg; + const int reps = std::max(1, (int)(2e10 / (2.0 * s.M * s.N * s.K))); + std::vector ds(std::min(reps, 200), d); + cx.warm(ds, 0.2); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds) / ds.size()); + const double m = median(t); + std::printf("%-26s %-36s %8.1f us %6.2f TFLOP/s\n", k.name.c_str(), s.name, m * 1e6, + 2.0 * s.M * s.N * s.K / m / 1e12); + } + } + return 0; +} diff --git a/edge0-strata/metal/tests/mma_parity.cpp b/edge0-strata/metal/tests/mma_parity.cpp new file mode 100644 index 00000000..3181dbad --- /dev/null +++ b/edge0-strata/metal/tests/mma_parity.cpp @@ -0,0 +1,178 @@ +// metal/tests/mma_parity.cpp - the simdgroup-matrix GEMV (mma.metal): parity vs double for every dense format, each +// token of a window bitwise a one-token call, and its speed on the packs' shapes against the per-lane GEMV. +// +// mma_parity parity +// mma_parity --bench speed, 1 / 3 / 8 tokens +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include + +namespace { +struct MmaArgs { uint32_t n_in, n_out, row_bytes, x_stride, y_stride; }; +struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +int xkind(int t) { + switch (t) { + case ref::Q2_0: return 1; + case ref::Q4_K: case ref::Q5_K: return 2; + case ref::Q6_K: case ref::Q3_K: return 3; + default: return 0; + } +} +} // namespace + +int main(int argc, char** argv) { + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "gemv.metal", "mma.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); + const int MAXT = 8; + mtl::Pipeline pprep = cx.pipeline(lib, "xprep"); + const std::vector types = {ref::BF16, ref::Q8_0, ref::Q4_0, ref::Q5_0, ref::IQ4_NL, ref::IQ4_XS, ref::Q2_0, + ref::IQ3_S, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K}; + if (!bench) { + int fails = 0; + for (int t : types) + for (int n_in : {2560, 640, 6144}) { + if (n_in % ref::block_elems(t) || (n_in == 640 && ref::block_elems(t) == 256)) continue; + const int n_out = 100; + const size_t rb = ref::row_bytes(t, n_in); + mtl::Buffer W = cx.alloc(rb * n_out), x = cx.alloc(4 * n_in * MAXT), y = cx.alloc(4 * n_out * MAXT), + y1 = cx.alloc(4 * n_out * MAXT); + ref::random_blocks(t, W.as(), n_out, n_in, 41 + t, 0.02f); + ref::Rng r(9); + for (int i = 0; i < n_in * MAXT; ++i) x.as()[i] = r.sym(); + std::vector wf((size_t)n_in * n_out); + ref::dequant(t, W.ptr, wf.data(), wf.size()); + for (int nt : {1, 3, 8}) { + mtl::Pipeline p = cx.pipeline(lib, "gemv_mma_" + std::to_string(t) + "_2", mtl::Constants().i(4, nt)); + MmaArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)rb, (uint32_t)n_in, (uint32_t)n_out}; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = (n_out + 15) / 16; + d.tg[0] = 64; + cx.run({d}); + double se = 0, sr = 0; + for (int tt = 0; tt < nt; ++tt) + for (int o = 0; o < n_out; ++o) { + double ref = 0; + for (int k = 0; k < n_in; ++k) ref += (double)wf[(size_t)o * n_in + k] * x.as()[tt * n_in + k]; + const double e = y.as()[tt * n_out + o] - ref; + se += e * e; + sr += ref * ref; + } + // the window property: token tt's column vs a one-token call on x shifted to row 0 + int notex = 0; + if (nt > 1) { + mtl::Pipeline p1 = cx.pipeline(lib, "gemv_mma_" + std::to_string(t) + "_2", mtl::Constants().i(4, 1)); + for (int tt = 0; tt < nt; ++tt) { + mtl::Dispatch d1 = d; + d1.p = &p1; + d1.args = {mtl::Arg::B(W), mtl::Arg::B(x, 4 * (size_t)tt * n_in), mtl::Arg::B(y1), mtl::Arg::V(a)}; + // rows past the window come from the following tokens' memory: padding only + cx.run({d1}); + notex += std::memcmp(y1.ptr, y.as() + tt * n_out, 4 * n_out) != 0; + } + } + const double rel = std::sqrt(se / sr); + const bool ok = rel < 1e-5 && notex == 0; + fails += !ok; + std::printf("mma parity %-7s n_in %5d ntok %d: rel L2 %.2e, window == one token: %s %s\n", + ref::type_name(t), n_in, nt, rel, nt == 1 ? "-" : notex ? "NO" : "bitwise", ok ? "PASS" : "FAIL"); + } + cx.free(W); + cx.free(x); + cx.free(y); + cx.free(y1); + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; + } + struct Shape { const char* name; int type, n_in, n_out; }; + const std::vector shapes = {{"gdn qkv", ref::IQ4_XS, 2560, 10240}, {"gdn qkv", ref::Q3_K, 2560, 10240}, + {"gdn qkv", ref::Q4_K, 2560, 10240}, {"gdn qkv", ref::Q6_K, 2560, 10240}, + {"gdn out", ref::Q5_K, 6144, 2560}, {"qsa q", ref::Q2_0, 2560, 12288}, + {"router", ref::BF16, 2560, 512}, {"lm head", ref::Q5_K, 2560, 248320}}; + std::printf("%-9s %-7s %5s | %9s %9s | %s\n", "shape", "type", "ntok", "lane us", "mma us", "mma GB/s (best SG)"); + for (const Shape& s : shapes) { + const size_t rb = ref::row_bytes(s.type, s.n_in), mb = rb * s.n_out; + const int copies = (int)std::max(4, std::min(64, (1ull << 30) / mb)); + mtl::Buffer W = cx.alloc(mb * copies), x = cx.alloc(4 * s.n_in * MAXT), y = cx.alloc(4 * (size_t)s.n_out * MAXT), + xp = cx.alloc(4 * (4 * s.n_in * MAXT + s.n_in / 16 * MAXT)); + { + std::vector tmp(mb); + ref::random_blocks(s.type, tmp.data(), s.n_out, s.n_in, 3, 0.02f); + for (int c = 0; c < copies; ++c) std::memcpy(W.as() + (size_t)c * mb, tmp.data(), mb); + } + ref::Rng r(5); + for (int i = 0; i < s.n_in * MAXT; ++i) x.as()[i] = r.sym(); + { + mtl::Dispatch d; + d.p = &pprep; + const uint32_t n = s.n_in, nt = MAXT; + d.args = {mtl::Arg::B(x), mtl::Arg::B(xp), mtl::Arg::V(n), mtl::Arg::V(nt)}; + d.tg[0] = 256; + d.grid[0] = (s.n_in / 16 * MAXT + 255) / 256; + cx.run({d}); + } + const size_t plane = (size_t)4 * s.n_in * MAXT; + for (int nt : {1, 3, 8}) { + auto time = [&](std::vector& ds) { + cx.warm(ds, 0.1); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds) / ds.size()); + return median(t); + }; + // the per-lane GEMV with its usual tile + const int L = s.n_in >= 6144 ? 32 : (s.n_out <= 640 || s.type == ref::Q2_0 || s.type == ref::BF16) ? 32 : 16; + const int S = s.n_in >= 6144 ? 8 : 4, K = s.n_in >= 6144 ? 2 : 1; + mtl::Pipeline pl = cx.pipeline(lib, "gemv_" + std::to_string(s.type) + "_" + std::to_string(L) + "_1_" + + std::to_string(S) + "_" + std::to_string(K), mtl::Constants().i(0, nt)); + GemvArgs ga{(uint32_t)s.n_in, (uint32_t)s.n_out, (uint32_t)nt, (uint32_t)rb, (uint32_t)s.n_in, (uint32_t)s.n_out}; + std::vector dl; + for (int k = 0; k < std::max(copies, 16); ++k) { + mtl::Dispatch d; + d.p = &pl; + d.args = {mtl::Arg::B(W, (size_t)(k % copies) * mb), mtl::Arg::B(xp, plane * xkind(s.type)), + mtl::Arg::B(xp, plane * 4), mtl::Arg::B(y), mtl::Arg::V(ga)}; + const int rpt = (S / K) * (32 / L); + d.grid[0] = (s.n_out + rpt - 1) / rpt; + d.tg[0] = S * 32; + dl.push_back(d); + } + const double tl = time(dl); + double best = 1e9; + int bsg = 0; + for (int sgn : {1, 2, 4}) { + mtl::Pipeline pm = cx.pipeline(lib, "gemv_mma_" + std::to_string(s.type) + "_" + std::to_string(sgn), + mtl::Constants().i(4, nt)); + MmaArgs ma{(uint32_t)s.n_in, (uint32_t)s.n_out, (uint32_t)rb, (uint32_t)s.n_in, (uint32_t)s.n_out}; + std::vector dm; + for (int k = 0; k < std::max(copies, 16); ++k) { + mtl::Dispatch d; + d.p = ± + d.args = {mtl::Arg::B(W, (size_t)(k % copies) * mb), mtl::Arg::B(x), mtl::Arg::B(y), mtl::Arg::V(ma)}; + d.grid[0] = (s.n_out + 8 * sgn - 1) / (8 * sgn); + d.tg[0] = 32 * sgn; + dm.push_back(d); + } + const double tm = time(dm); + if (tm < best) { best = tm; bsg = sgn; } + } + std::printf("%-9s %-7s %5d | %9.2f %9.2f | %6.1f (SG %d) lane %6.1f GB/s\n", s.name, ref::type_name(s.type), nt, + tl * 1e6, best * 1e6, mb / best / 1e9, bsg, mb / tl / 1e9); + } + cx.free(W); + cx.free(x); + cx.free(y); + cx.free(xp); + } + return 0; +} diff --git a/edge0-strata/metal/tests/moe_mm_parity.cpp b/edge0-strata/metal/tests/moe_mm_parity.cpp new file mode 100644 index 00000000..dd04ba39 --- /dev/null +++ b/edge0-strata/metal/tests/moe_mm_parity.cpp @@ -0,0 +1,268 @@ +// metal/tests/moe_mm_parity.cpp - the prompt-chunk routed experts (moe_mm.metal): plan, gate/up GEMM, down GEMM, +// combine, against a CPU reference that rounds the activations and weights to fp16 as the kernels do (the matrix +// units' operands) - and against the exact float reference for the size of that rounding; then the speed. +// +// moe_mm_parity [--gu Q2_0] [--dn Q2_0] [--T 256] [--tile 32,64,64,4] parity (64 experts) +// moe_mm_parity --bench [--gu ..] [--dn ..] [--T 2048] [--tile ..] per-layer time, 512 experts +#include "coop_probe.h" +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +uint32_t N_EMBD = 2560; +constexpr int N_FF = 640, TOPK = 10; +struct MmPlanArgs { uint32_t n_ent, topk, n_expert, tm, n_blocks; }; +struct MmArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +int type_of(const std::string& s) { + for (int t : {ref::TQ2_T, ref::TQK6, ref::TQK7, ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::Q8_0, ref::Q5_1}) + if (s == ref::type_name(t)) return t; + std::exit(2); +} +float h16(float v) { return ref::f16_to_f32(ref::f32_to_f16(v)); } +} // namespace + +int main(int argc, char** argv) { + int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 64, SG = 2, DTN = -1, DTK = -1, DSG = -1; + bool bench = false, v2 = false, v3 = false, v4 = false, v5 = false, stage_gu = false; + int MS = 1; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--gu") gu = type_of(argv[++i]); + else if (a == "--dn") dn = type_of(argv[++i]); + else if (a == "--T") T = std::atoi(argv[++i]); + else if (a == "--latent") N_EMBD = 1280; + else if (a == "--latent1792") N_EMBD = 1792; + else if (a == "--stage-gu") stage_gu = true; + else if (a == "--tile") std::sscanf(argv[++i], "%d,%d,%d,%d", &TM, &TN, &TK, &SG); + else if (a == "--dtile") std::sscanf(argv[++i], "%d,%d,%d", &DTN, &DTK, &DSG); // the down kernel's TN,TK,SG + else if (a == "--bench") bench = true; + else if (a == "--v2") v2 = true; + else if (a == "--v3") v2 = v3 = true; // v2's gather + the register-operand GEMMs (TN per simdgroup) + else if (a == "--v4") { v2 = v3 = v4 = true; MS = std::atoi(argv[++i]); } // mm_gu4 with MS 16-row subtiles + else if (a == "--v5") { v2 = v3 = true; v5 = true; } // mm_gu5 (weight blocks staged in threadgroup memory) + } + if (bench && T == 256) T = 2048; + if (v4) TM = 16 * MS; // the plan's tile height; mm_gu4's name carries MS instead + if (stage_gu && (bench || !v2 || v3 || gu != ref::Q3_K || TM != 32 || TN != 64 || TK != 32 || SG != 2)) { + std::fprintf(stderr, "--stage-gu requires the Q3_K v2 parity case with tile 32,64,32,2\n"); + return 2; + } + mtl::Context cx; + std::vector sources{"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}; + if (stage_gu) sources.push_back("latent_prompt_q3_stage.metal"); + const int lib = cx.compile(sources); + const int NE = bench ? 512 : 64, LAYERS = bench ? 4 : 1; + const size_t gu_row = ref::row_bytes(gu, N_EMBD), d_row = ref::row_bytes(dn, N_FF); + const size_t up_off = (size_t)N_FF * gu_row, down_off = 2 * up_off, blob = down_off + (size_t)N_EMBD * d_row; + std::printf("device %s; %s/%s, T %d, %d experts, tile TM %d TN %d TK %d SG %d\n", cx.device_name().c_str(), + ref::type_name(gu), ref::type_name(dn), T, NE, TM, TN, TK, SG); + mtl::Buffer W = cx.alloc(blob * NE * LAYERS); + { + std::vector tmp(blob * 16); + for (int e = 0; e < 16; ++e) { + ref::random_blocks(gu, tmp.data() + e * blob, 2 * N_FF, N_EMBD, 100 + e, 0.02f); + ref::random_blocks(dn, tmp.data() + e * blob + down_off, N_EMBD, N_FF, 200 + e, 0.02f); + } + for (size_t e = 0; e < (size_t)NE * LAYERS; ++e) std::memcpy(W.as() + e * blob, tmp.data() + (e % 16) * blob, blob); + } + const int nent = T * TOPK, nblk = (nent + 1023) / 1024; + const int maxrows = nent + NE * TM; + mtl::Buffer x = cx.alloc(4 * (size_t)T * N_EMBD), ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent), + bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE), + tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows), + sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent), H = cx.alloc(4 * (size_t)maxrows * N_FF), + YE = cx.alloc(4 * (size_t)maxrows * N_EMBD), y = cx.alloc(4 * (size_t)T * N_EMBD), + XS = cx.alloc(2 * (size_t)maxrows * N_EMBD), HS = cx.alloc(2 * (size_t)maxrows * N_FF); + ref::Rng r(7); + for (size_t i = 0; i < (size_t)T * N_EMBD; ++i) x.as()[i] = r.sym(); + for (int t = 0; t < T; ++t) { + std::set used; + float s = 0; + for (int k = 0; k < TOPK; ++k) { + int e; + do e = (int)(r.next() % NE); while (used.count(e)); + used.insert(e); + ids.as()[t * TOPK + k] = e; + s += (w.as()[t * TOPK + k] = 0.2f + r.uni()); + } + for (int k = 0; k < TOPK; ++k) w.as()[t * TOPK + k] /= s; + } + mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"), + pclear = cx.pipeline(lib, "mm_clear"), pcomb = cx.pipeline(lib, "mm_combine"); + if (DTN < 0) { DTN = TN; DTK = TK; DSG = SG; } + if (dn == ref::IQ4_NL && DTK < 32) DTK = 32; + const std::string tl = std::to_string(TM) + "_" + std::to_string(TN) + "_" + std::to_string(TK) + "_" + std::to_string(SG); + const std::string dtl = std::to_string(TM) + "_" + std::to_string(DTN) + "_" + std::to_string(DTK) + "_" + std::to_string(DSG); + mtl::Constants probe; + probe.i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0); + const bool k4 = coop_k4(cx); + if (v3) std::printf("right-input layout: %s\n", k4 ? "runs of 4 k" : "other (per-element decode)"); + mtl::Buffer w16; + mtl::Pipeline pstage; + if (stage_gu) { + w16 = cx.alloc(2ull * NE * 2 * N_FF * N_EMBD); + pstage = cx.pipeline(lib, "latent_prompt_q3_stage"); + } + mtl::Pipeline pgu = stage_gu ? cx.pipeline(lib, "latent_prompt_gu_f16") + : v5 ? cx.pipeline(lib, "mm_gu5_" + std::to_string(gu) + "_" + std::to_string(TN) + "_" + std::to_string(SG), + mtl::Constants().b(7, k4)) + : v4 ? cx.pipeline(lib, "mm_gu4_" + std::to_string(gu) + "_" + std::to_string(MS) + tl.substr(tl.find('_')), + mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0)) + : v3 ? cx.pipeline(lib, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr)) + : cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe), + pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4)) + : cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)), + pgat = cx.pipeline(lib, "mm_gather"); + MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk}; + auto encode = [&](int layer, std::vector& ds, bool plan, bool gemm) { + if (plan) { + const uint32_t nclr = maxrows; + mtl::Dispatch c; c.p = &pclear; c.args = {mtl::Arg::B(stok), mtl::Arg::V(nclr)}; c.grid[0] = (maxrows + 255) / 256; c.tg[0] = 256; ds.push_back(c); + mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds.push_back(h); + meta.as()[2] = N_FF / TN; + mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa)}; sc.tg[0] = NE; ds.push_back(sc); + mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds.push_back(pl); + } + if (gemm) { + MmArgs ma{(uint64_t)layer * NE * blob, (uint64_t)layer * NE * blob + up_off, (uint64_t)layer * NE * blob + down_off, + blob, blob, (uint32_t)gu_row, (uint32_t)d_row, N_EMBD, N_FF}; + if (v2 || v3) { // v3 is v2's gather plus the register-operand GEMMs + const uint32_t ne = N_EMBD; + mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(x), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; + ga.grid[0] = meta.as()[1]; ga.tg[0] = 256; ds.push_back(ga); + } + if (stage_gu) { + const uint32_t q3[4] = {N_EMBD, N_FF, uint32_t(gu_row), uint32_t(blob)}; + mtl::Dispatch stage; + stage.p = &pstage; + stage.args = {mtl::Arg::B(W), mtl::Arg::B(w16), mtl::Arg::V(q3)}; + stage.grid[0] = (2 * N_FF * (N_EMBD / 32) + 127) / 128; + stage.grid[1] = NE; + stage.tg[0] = 128; + ds.push_back(stage); + } + mtl::Dispatch g; g.p = &pgu; + if (v2 || v3) g.args = {mtl::Arg::B(stage_gu ? w16 : W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + else g.args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(H), mtl::Arg::B(stok), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + g.grid[0] = N_FF / (v3 ? TN * SG : TN); g.grid[1] = meta.as()[0]; g.tg[0] = 32 * SG; ds.push_back(g); + mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 || v3 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as()[0]; d.tg[0] = 32 * ((v2 || v3) ? DSG : SG); ds.push_back(d); + const uint32_t ca[4] = {(uint32_t)T, N_EMBD, TOPK, 0}; + mtl::Dispatch cb; cb.p = &pcomb; cb.args = {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(y), mtl::Arg::V(ca)}; + cb.grid[0] = (T * N_EMBD + 255) / 256; cb.tg[0] = 256; ds.push_back(cb); + } + }; + { + std::vector ds; + encode(0, ds, true, false); + cx.run(ds); // the plan first: the GEMMs' grid is the tile count + } + const uint32_t ntiles = meta.as()[0]; + std::printf("plan: %u tiles of %d, %u padded rows for %d entries\n", ntiles, TM, meta.as()[1], nent); + if (!bench) { + // the plan: every entry placed once, rows of an expert contiguous and in entry order + bool ok_plan = true; + std::vector seen(maxrows, 0); + for (int i = 0; i < nent; ++i) { + const int row = pos.as()[i]; + if (row < 0 || row >= (int)meta.as()[1] || stok.as()[row] != i / TOPK) ok_plan = false; + else seen[row]++; + } + for (int i = 0; i < maxrows; ++i) if (seen[i] > 1) ok_plan = false; + std::vector ds; + encode(0, ds, false, true); + cx.run(ds); + // reference: fp16 operands (x, weights, h), fp32+ accumulation; and the exact one + double se16 = 0, se = 0, sr = 0; + std::vector wg((size_t)N_FF * N_EMBD), wu(wg.size()), wd((size_t)N_EMBD * N_FF); + std::vector yr16((size_t)T * N_EMBD, 0.0), yr((size_t)T * N_EMBD, 0.0); + for (int e = 0; e < NE; ++e) { + const uint8_t* b = W.as() + e * blob; + bool any = false; + for (int i = 0; i < nent && !any; ++i) any = ids.as()[i] == e; + if (!any) continue; + ref::dequant(gu, b, wg.data(), wg.size()); + ref::dequant(gu, b + up_off, wu.data(), wu.size()); + ref::dequant(dn, b + down_off, wd.data(), wd.size()); + for (int i = 0; i < nent; ++i) { + if (ids.as()[i] != e) continue; + const int t = i / TOPK; + const float* xt = x.as() + (size_t)t * N_EMBD; + std::vector h16v(N_FF), hv(N_FF); + for (int f = 0; f < N_FF; ++f) { + double g = 0, u = 0, g16 = 0, u16 = 0; + for (int k = 0; k < N_EMBD; ++k) { + const float xk = xt[k], xh = h16(xk); + g += (double)wg[(size_t)f * N_EMBD + k] * xk; + u += (double)wu[(size_t)f * N_EMBD + k] * xk; + g16 += (double)h16(wg[(size_t)f * N_EMBD + k]) * xh; + u16 += (double)h16(wu[(size_t)f * N_EMBD + k]) * xh; + } + hv[f] = g / (1 + std::exp(-g)) * u; + h16v[f] = h16((float)(g16 / (1 + std::exp(-g16)) * u16)); + } + const double wt = w.as()[i]; + for (int d = 0; d < N_EMBD; ++d) { + double s = 0, s16 = 0; + for (int f = 0; f < N_FF; ++f) { + s += (double)wd[(size_t)d * N_FF + f] * hv[f]; + s16 += (double)h16(wd[(size_t)d * N_FF + f]) * h16v[f]; + } + yr[(size_t)t * N_EMBD + d] += wt * s; + yr16[(size_t)t * N_EMBD + d] += wt * s16; + } + } + } + for (size_t i = 0; i < yr.size(); ++i) { + const double g = y.as()[i]; + se16 += (g - yr16[i]) * (g - yr16[i]); + se += (g - yr[i]) * (g - yr[i]); + sr += yr[i] * yr[i]; + } + const double rel16 = std::sqrt(se16 / sr), rel = std::sqrt(se / sr); + const bool ok = ok_plan && rel16 < 2e-4; + std::printf("parity: plan %s; y vs fp16-operand reference rel L2 %.2e, vs exact float reference %.2e %s\n", + ok_plan ? "ok" : "WRONG", rel16, rel, ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; + } + // bench: the plan, the two GEMMs and the combine, LAYERS distinct weight sets rotated + double flops = 0; + { + const double rows = meta.as()[1]; + flops = 2.0 * nent * (2.0 * N_FF * N_EMBD + (double)N_EMBD * N_FF); + std::printf("useful GFLOP per layer %.1f (padded rows %.0f = %.2fx entries)\n", flops / 1e9, rows, rows / nent); + } + auto timeit = [&](bool plan, bool gemm, int part) { + std::vector ds; + for (int l = 0; l < 8; ++l) { + std::vector one; + encode(l % LAYERS, one, plan, gemm); + if (part >= 0 && gemm) { ds.push_back(one[part]); continue; } + for (auto& d : one) ds.push_back(d); + } + cx.warm(ds, 0.3); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds) / 8); + return median(t); + }; + const int o = v2 ? 1 : 0; + const double tp = timeit(true, false, -1), tg = timeit(false, true, o), td = timeit(false, true, o + 1), + tc = timeit(false, true, o + 2), ta = timeit(false, true, -1); + if (v2) std::printf("gather %.1f us\n", timeit(false, true, 0) * 1e6); + const double fgu = 2.0 * nent * 2.0 * N_FF * N_EMBD, fdn = 2.0 * nent * N_EMBD * N_FF; + std::printf("per layer: plan %.1f us, gate/up %.1f us (%.1f TFLOP/s useful), down %.1f us (%.1f TFLOP/s), combine %.1f us;" + " GEMMs + combine %.1f us = %.1f TFLOP/s; %.1f us per token per layer -> 48 layers: %.0f tok/s (experts only)\n", + tp * 1e6, tg * 1e6, fgu / tg / 1e12, td * 1e6, fdn / td / 1e12, tc * 1e6, ta * 1e6, flops / ta / 1e12, + (ta + tp) / T * 1e6, T / (48 * (ta + tp))); + return 0; +} diff --git a/edge0-strata/metal/tests/pair_parity.cpp b/edge0-strata/metal/tests/pair_parity.cpp new file mode 100644 index 00000000..1e40d19f --- /dev/null +++ b/edge0-strata/metal/tests/pair_parity.cpp @@ -0,0 +1,118 @@ +// metal/tests/pair_parity.cpp - gemv_pair (the shared expert's gate + up + SwiGLU in one dispatch) vs double, for +// format pairs of the packs, 1 / 3 / 8 tokens; and the raw-activation Q2_0 GEMV (type 142) the down then uses. +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include + +namespace { +struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; +int xkind(int t) { + switch (t) { + case ref::Q2_0: return 1; + case ref::Q4_K: case ref::Q5_K: return 2; + case ref::Q6_K: case ref::Q3_K: return 3; + default: return 0; + } +} +} // namespace + +int main() { + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "gemv.metal"}); + mtl::Pipeline pprep = cx.pipeline(lib, "xprep"); + const int NIN = 2560, NOUT = 640, MAXT = 8; + mtl::Buffer x = cx.alloc(4 * NIN * MAXT), xp = cx.alloc(4 * (4 * NIN * MAXT + NIN / 16 * MAXT)), h = cx.alloc(4 * NOUT * MAXT); + ref::Rng r(4); + for (int i = 0; i < NIN * MAXT; ++i) x.as()[i] = r.sym(); + { + mtl::Dispatch d; + d.p = &pprep; + const uint32_t n = NIN, nt = MAXT; + d.args = {mtl::Arg::B(x), mtl::Arg::B(xp), mtl::Arg::V(n), mtl::Arg::V(nt)}; + d.tg[0] = 256; + d.grid[0] = (NIN / 16 * MAXT + 255) / 256; + cx.run({d}); + } + const size_t plane = (size_t)4 * NIN * MAXT; + const std::vector> pairs = {{ref::Q2_0, ref::Q3_K}, {ref::Q3_K, ref::Q2_0}, {ref::IQ4_XS, ref::IQ3_S}, + {ref::Q4_K, ref::Q6_K}, {ref::Q5_K, ref::Q4_K}, {ref::Q6_K, ref::IQ3_S}, + {ref::Q2_0, ref::Q2_0}, {ref::IQ3_S, ref::Q4_K}}; + int fails = 0; + for (auto [fg, fu] : pairs) { + const size_t rg = ref::row_bytes(fg, NIN), ru = ref::row_bytes(fu, NIN); + mtl::Buffer Wg = cx.alloc(rg * NOUT), Wu = cx.alloc(ru * NOUT); + ref::random_blocks(fg, Wg.as(), NOUT, NIN, 7 + fg, 0.02f); + ref::random_blocks(fu, Wu.as(), NOUT, NIN, 9 + fu, 0.02f); + std::vector wg((size_t)NIN * NOUT), wu((size_t)NIN * NOUT); + ref::dequant(fg, Wg.ptr, wg.data(), wg.size()); + ref::dequant(fu, Wu.ptr, wu.data(), wu.size()); + for (int nt : {1, 3, 8}) { + mtl::Pipeline p = cx.pipeline(lib, "gemv_pair_" + std::to_string(fg) + "_" + std::to_string(fu), + mtl::Constants().i(0, nt)); + GemvArgs a{(uint32_t)NIN, (uint32_t)NOUT, (uint32_t)nt, (uint32_t)rg, (uint32_t)NIN, (uint32_t)NOUT}; + const uint32_t rbu = (uint32_t)ru; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(Wg), mtl::Arg::B(xp, plane * xkind(fg)), mtl::Arg::B(Wu), mtl::Arg::B(xp, plane * xkind(fu)), + mtl::Arg::B(xp, plane * 4), mtl::Arg::B(h), mtl::Arg::V(a), mtl::Arg::V(rbu)}; + d.grid[0] = NOUT / 4; + d.tg[0] = 128; + cx.run({d}); + double se = 0, sr = 0; + for (int t = 0; t < nt; ++t) + for (int o = 0; o < NOUT; ++o) { + double g = 0, u = 0; + for (int k = 0; k < NIN; ++k) { + g += (double)wg[(size_t)o * NIN + k] * x.as()[t * NIN + k]; + u += (double)wu[(size_t)o * NIN + k] * x.as()[t * NIN + k]; + } + const double ref = g / (1.0 + std::exp(-g)) * u; + const double e = h.as()[t * NOUT + o] - ref; + se += e * e; + sr += ref * ref; + } + const double rel = std::sqrt(se / sr); + fails += !(rel < 1e-5); + std::printf("pair %-7s + %-7s ntok %d: rel L2 %.2e %s\n", ref::type_name(fg), ref::type_name(fu), nt, rel, + rel < 1e-5 ? "PASS" : "FAIL"); + } + cx.free(Wg); + cx.free(Wu); + } + // the raw-activation Q2_0 GEMV + { + const int ni = 640, no = 2560; + const size_t rb = ref::row_bytes(ref::Q2_0, ni); + mtl::Buffer W = cx.alloc(rb * no), y = cx.alloc(4 * no * MAXT); + ref::random_blocks(ref::Q2_0, W.as(), no, ni, 3, 0.02f); + std::vector wf((size_t)ni * no); + ref::dequant(ref::Q2_0, W.ptr, wf.data(), wf.size()); + mtl::Pipeline p = cx.pipeline(lib, "gemv_142_8_1_4_1", mtl::Constants().i(0, 3)); + GemvArgs a{(uint32_t)ni, (uint32_t)no, 3, (uint32_t)rb, (uint32_t)ni, (uint32_t)no}; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(x), mtl::Arg::B(y), mtl::Arg::V(a)}; + d.grid[0] = no / 16; + d.tg[0] = 128; + cx.run({d}); + double se = 0, sr = 0; + for (int t = 0; t < 3; ++t) + for (int o = 0; o < no; ++o) { + double ref = 0; + for (int k = 0; k < ni; ++k) ref += (double)wf[(size_t)o * ni + k] * x.as()[t * ni + k]; + const double e = y.as()[t * no + o] - ref; + se += e * e; + sr += ref * ref; + } + const double rel = std::sqrt(se / sr); + fails += !(rel < 1e-5); + std::printf("Q2_0 raw-x GEMV 640 -> 2560, 3 tokens: rel L2 %.2e %s\n", rel, rel < 1e-5 ? "PASS" : "FAIL"); + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/prompt_chain.cpp b/edge0-strata/metal/tests/prompt_chain.cpp new file mode 100644 index 00000000..876afeca --- /dev/null +++ b/edge0-strata/metal/tests/prompt_chain.cpp @@ -0,0 +1,555 @@ +// metal/tests/prompt_chain.cpp - a whole PROMPT through every layer on the REAL weights of a GGUF (mmapped, wrapped +// no-copy), in chunks of C tokens (one command buffer a chunk): the GPU time of prompt processing with the prompt +// kernels of metal/. A timing of the real byte and FLOP stream, not a numerical check: the tokens are random ids +// (their real embedding rows start every hyper-connection stream), the parity programs check the kernels. +// +// prompt_chain [--prompt P,P,..] [--chunk C] [--breakdown] [--only CLASS] [--skip CLASS] +// +// Per layer and chunk (T tokens; activations fp32 rows, fp16 copies as the matrix units' operands): +// HC read gr_norm_h (fp16 xn) -> GEMM [w_down; w_inject] -> gr_epi_down -> gemm_hc_up (the GEMM with the +// mixing epilogue); PC_HC=plain: gr_norm -> act_f16 -> GEMM -> gr_epi_down -> GEMM -> gr_epi_up +// GDN GEMMs qkv / gate / alpha / beta -> gdn_pre (parallel) -> gdn_scan_r (the state scan) -> gdn_post +// (output norm and gate, fp16) -> GEMM out; PC_GDN=step: the decode step over the T tokens instead +// QSA GEMMs q / k / v / indexer q, k -> qsa_prep (T tokens) -> score + top-k (past the budget) -> attn -> +// merge -> GEMM out +// MoE GEMM router -> route_topk -> plan (hist, scan, clear, place) -> gather -> mm_gu2 -> mm_down2 -> +// combine; shared expert GEMMs gate, up -> swiglu_h -> GEMM down; GEMM shared gate; moe_out +// Every dense matrix is dequantized to fp16 (dq_f16) once a chunk, before its GEMM. After the last chunk: the final +// HC read and the LM head on the LAST row only (the decode kernels: gr_norm / gr_down / gr_up, GEMV, argmax). +// Left out: the PLE path and the MTP draft. +#include "coop_probe.h" +#include "gguf_lite.h" +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int N = 2560, HC = 4, LR = 320, NFF = 640, TOPK = 10, NEXP = 512, NV = 248320, QB = 2051, QCHUNKS = 33; +struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; +struct EpiArgs { uint32_t n_tok, n_embd, hc, hc_lr; }; +struct DqArgs { uint32_t n_in, n_out, row_bytes, perm; }; +struct GemmArgs { int M, N, K, pad; }; +struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; +struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; }; +struct GdnPArgs { uint32_t n_tok, h_k, h_v, channels; float eps; uint32_t f32; }; +struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; }; +struct MmPlanArgs { uint32_t n_ent, topk, n_expert, tm, n_blocks; }; +struct MmArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff; }; +struct ArgArgs { uint32_t n_vocab, parts; }; + +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +bool fmt_unit32(int t) { return t == ref::IQ4_XS || t == ref::IQ4_NL; } // mm_gu2 needs a K chunk >= 32 +int unit_of(int t) { return t == ref::Q2_0 || t == ref::IQ3_S || t == ref::BF16 ? 16 : 32; } + +} // namespace + +int main(int argc, char** argv) { + if (argc < 2) { std::fprintf(stderr, "usage: prompt_chain [--prompt P,..] [--chunk C] [--breakdown]\n"); return 2; } + std::vector prompts = {1024, 4096, 32768}; + int C = 2048, reps = 2; + bool breakdown = false; + std::string only, skip; + for (int i = 2; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--prompt") { + prompts.clear(); + for (char* s = argv[++i]; *s;) { + prompts.push_back((int)std::strtol(s, &s, 10)); + if (*s == ',') ++s; + } + } else if (a == "--chunk") C = std::atoi(argv[++i]); + else if (a == "--reps") reps = std::atoi(argv[++i]); + else if (a == "--breakdown") breakdown = true; + else if (a == "--only") only = argv[++i]; + else if (a == "--skip") skip = argv[++i]; + } + const int PMAX = *std::max_element(prompts.begin(), prompts.end()); + gguf::File f; + std::string err; + if (!f.open(argv[1], err)) { std::fprintf(stderr, "%s\n", err.c_str()); return 2; } + mtl::Context cx; + std::printf("device %s; %s: %zu tensors, %.1f GB; chunk %d\n", cx.device_name().c_str(), argv[1], f.tensors.size(), + f.file_bytes / 1e9, C); + { + const double t0 = mtl::now_s(); + madvise(f.map, f.map_bytes, MADV_WILLNEED); + volatile uint64_t s = 0; + for (size_t o = 0; o < f.file_bytes; o += 16384) s += ((const uint8_t*)f.map)[o]; + std::printf("file touched in %.1f s\n", mtl::now_s() - t0); + } + mtl::Buffer W = cx.wrap(f.map, f.map_bytes); + auto T_ = [&](const std::string& name) -> const gguf::Tensor& { + const gguf::Tensor* t = f.find(name); + if (!t) { std::fprintf(stderr, "missing tensor %s\n", name.c_str()); std::exit(2); } + return *t; + }; + const double tc0 = mtl::now_s(); + const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", + "gemv.metal", "prompt.metal"}); + const int lgdn = cx.compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const int lqsa = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + const int lsmp = cx.compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("kernels compiled in %.1f s\n", mtl::now_s() - tc0); + std::map pcache; + const bool k4 = coop_k4(cx); + std::printf("matmul2d right-input layout: %s\n", k4 ? "runs of 4 k (per-lane run decoders)" : "other (per-element decode)"); + auto P = [&](int l, const std::string& fn, const mtl::Constants& c = {}, const std::string& key = "") -> const mtl::Pipeline* { + const std::string k = key.empty() ? fn : key; + auto it = pcache.find(k); + if (it == pcache.end()) it = pcache.emplace(k, cx.pipeline(l, fn, c)).first; + return &it->second; + }; + + // ---- buffers (MT = the chunk's rows) + const size_t MT = C; + const int MAXC = (PMAX + 64 + 63) / 64 * 64, MB = MAXC / 4 + 2; + const int TM = 16; + const int maxrows = (int)MT * TOPK + NEXP * TM, nblk_max = ((int)MT * TOPK + 1023) / 1024; + auto F32 = [&](size_t n) { return cx.alloc(4 * n); }; + auto F16 = [&](size_t n) { return cx.alloc(2 * n); }; + mtl::Buffer R = F32(MT * HC * N), bo = F32(MT * N), inj[2] = {F32(MT * HC), F32(MT * HC)}, xn = F32(MT * HC * N), + xnh = F16(MT * HC * N), V = F32(MT * (LR + HC)), loh = F16(MT * LR), G = F32(MT * HC * N), mixed = F32(MT * N), + mixh = F16(MT * N); + mtl::Buffer qkv = F32(MT * 10240), z = F32(MT * 6144), al = F32(MT * 48), be = F32(MT * 48), yg = F32(MT * 6144), + ygh = F16(MT * 6144), gqn = F32(MT * 16 * 128), gkn = F32(MT * 16 * 128), gvc = F32(MT * 6144), + ggb = F32(MT * 48 * 2), gkq = F32(MT * 16), go = F32(MT * 6144); + const char* hc_env = std::getenv("PC_HC"); + const std::string hc_mode = hc_env ? hc_env : "fused"; // fused | plain + const char* qsa_env = std::getenv("PC_QSA"); + // flash (one pass on the matrix units, no partials) | mm (split-K on the matrix units + merge) | v1 (qsa_attn, + // fp32 + merge). The M5: flash (= mm + merge or better); the M4: v1 attention (faster there) with the GEMM score. + const std::string qsa_mode = qsa_env ? qsa_env : (k4 ? "flash" : "v1"); + // the expert GEMMs: "v2" (threadgroup-staged weights) or "v3:TM,TN,TK,SG" (register operands, TN rows a simdgroup) + const char* gu_env = std::getenv("PC_GU"); + const char* dn_env = std::getenv("PC_DN"); + // defaults (M5 Max, moe_mm_parity --v3 --bench): gate/up 16,32,32,4 (every format), down 16,64,32,2 + // without the M5's right-input layout (the M4) the per-element v3 decode is slower than v2: v2 there + const std::string gu_mode = gu_env ? gu_env : (k4 ? "v3:auto" : "v2"), dn_mode = dn_env ? dn_env : (k4 ? "v3:16,64,32,2" : "v2"); + auto v3_tile = [](const std::string& m, int* t) { std::sscanf(m.c_str() + 3, "%d,%d,%d,%d", &t[0], &t[1], &t[2], &t[3]); }; + int gu3[4] = {16, 64, 16, 2}, dn3[4] = {16, 64, 16, 2}; + if (gu_mode.rfind("v3:", 0) == 0 && gu_mode != "v3:auto") v3_tile(gu_mode, gu3); + else if (gu_mode == "v3:auto") { gu3[0] = 16; gu3[1] = 32; gu3[2] = 32; gu3[3] = 4; } + auto tl3 = [](const int* t) { return std::to_string(t[0]) + "_" + std::to_string(t[1]) + "_" + std::to_string(t[2]) + "_" + std::to_string(t[3]); }; + auto gu3_for = [&](int type) -> std::string { + (void)type; // one tile for every format since the paired-run decoders (W8) + return tl3(gu3); + }; + if (dn_mode.rfind("v3:", 0) == 0) v3_tile(dn_mode, dn3); + const bool gu_v3 = gu_mode.rfind("v3:", 0) == 0, dn_v3 = dn_mode.rfind("v3:", 0) == 0; + const char* gdn_env = std::getenv("PC_GDN"); + const std::string gdn_mode = gdn_env ? gdn_env : "scanr_4"; + mtl::Buffer qfull = F32(MT * 12288), kcur = F32(MT * 512), vcur = F32(MT * 512), iqr = F32(MT * 512), ikr = F32(MT * 128), + qcur = F32(MT * 24 * 256), qcurh = F16(MT * 24 * 256), qidx = F32(MT * 4 * 128), scores = F32(MT * MB), qids = cx.alloc(4 * MT * QB), + part = F32(MT * QCHUNKS * 24 * 258), XQ = F32(4 * MT * 6144 + MT * 384), xqh = F16(MT * 6144); + mtl::Buffer logits = F32(MT * NEXP), ids = cx.alloc(4 * MT * TOPK), wts = F32(MT * TOPK), + bcnt = cx.alloc(4 * (size_t)nblk_max * NEXP), boff = cx.alloc(4 * (size_t)nblk_max * NEXP), eoff = cx.alloc(4 * NEXP), + tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows), + sw = F32(maxrows), pos = cx.alloc(4 * MT * TOPK), XS = F16((size_t)maxrows * N), HS = F16((size_t)maxrows * NFF), + YE = F32((size_t)maxrows * N), yr = F32(MT * N), gsh = F32(MT * NFF), ush = F32(MT * NFF), hh = F16(MT * NFF), + ys = F32(MT * N), sgv = F32(MT); + // fp16 weight scratch, one per role + mtl::Buffer WD = F16((size_t)(LR + HC) * HC * N), WU = F16((size_t)HC * N * LR), Wqkv = F16((size_t)10240 * N), + Wz = F16((size_t)6144 * N), Wal = F16((size_t)48 * N), Wbe = F16((size_t)48 * N), Wout = F16((size_t)N * 6144), + Wq = F16((size_t)12288 * N), Wk = F16((size_t)512 * N), Wv = F16((size_t)512 * N), Wiq = F16((size_t)512 * N), + Wik = F16((size_t)128 * N), Wr = F16((size_t)NEXP * N), Wsg = F16((size_t)N), Wshg = F16((size_t)NFF * N), + Wshu = F16((size_t)NFF * N), Wshd = F16((size_t)N * NFF); + // the last row's final read and head + mtl::Buffer xn1 = F32(HC * N), lo1 = F32(LR), inj1 = F32(HC), X1 = F32(4 * N + N / 16), lm = F32(NV), apv = F32(120), + api = cx.alloc(4 * 120), tokout = cx.alloc(16); + mtl::Buffer cs = cx.alloc(8 * (size_t)(MAXC + 16) * 32); + for (int p = 0; p < MAXC + 16; ++p) + for (int i = 0; i < 32; ++i) { + const double th = (double)p * std::pow(1e7, -2.0 * i / 64.0); + cs.as()[(p * 32 + i) * 2] = (float)std::cos(th); + cs.as()[(p * 32 + i) * 2 + 1] = (float)std::sin(th); + } + std::vector gst, ghist[2], kq, vq, ksc, vsc, tail, dead, pooled; + for (int l = 0; l < 48; ++l) { + if (l % 4 != 3) { + gst.push_back(F32(128 * 48 * 128)); + ghist[0].push_back(F32(10240 * 3)); + ghist[1].push_back(F32(10240 * 3)); + } else { + kq.push_back(cx.alloc((size_t)MAXC * 2 * 256)); + vq.push_back(cx.alloc((size_t)MAXC * 2 * 256)); + ksc.push_back(cx.alloc(2 * (size_t)MAXC * 2 * 4)); + vsc.push_back(cx.alloc(2 * (size_t)MAXC * 2 * 4)); + tail.push_back(F32(4 * 128)); + dead.push_back(F32(128)); + pooled.push_back(F32((size_t)MB * 128)); + } + } + std::vector gpar(36, 0); + // embeddings of random token ids (the real rows) + std::vector E(MT * N); + { + const gguf::Tensor& te = T_("token_embd.weight"); + const size_t rb = ref::row_bytes(te.type, N); + ref::Rng r(17); + for (size_t t = 0; t < MT; ++t) { + const uint64_t id = r.next() % 200000; + ref::dequant(te.type, (const uint8_t*)f.map + te.offset + id * rb, E.data() + t * N, N); + } + } + meta.as()[2] = gu_v3 ? NFF / (gu3[1] * gu3[3]) : NFF / 64; + meta.as()[5] = dn_v3 ? N / (dn3[1] * dn3[3]) : N / 64; + + // ---- the chunk's dispatches + std::vector ds; + std::vector cls; + auto push = [&](const mtl::Dispatch& d, const std::string& c) { + if (!only.empty() && c != only) return; + if (!skip.empty() && c == skip) return; + ds.push_back(d); + cls.push_back(c); + }; + auto D = [&](const mtl::Pipeline* p, std::vector a, uint32_t gx, uint32_t gy, uint32_t tg, const std::string& c, + uint32_t gz = 1) { + mtl::Dispatch d; + d.p = p; + d.args = std::move(a); + d.grid[0] = gx; + d.grid[1] = gy; + d.grid[2] = gz; + d.tg[0] = tg; + push(d, c); + }; + const char* gk_env = std::getenv("PC_GEMM"); + auto gemm = [&](const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Y, int M, int Nn, int K, const std::string& c) { + int tm = 64, tn = Nn <= 64 ? 32 : 64; + if (gk_env && Nn > 64) std::sscanf(gk_env, "%d,%d", &tm, &tn); + const std::string fn = "gemm_f16_" + std::to_string(tm) + "_" + std::to_string(tn) + "_4"; + GemmArgs g{M, Nn, K, 0}; + D(P(lib, fn), {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, (Nn + tn - 1) / tn, (M + tm - 1) / tm, + 128, c); + }; + auto dq = [&](const gguf::Tensor& t, const mtl::Buffer& out, size_t out_off = 0, uint32_t perm = 0) { + const int n_in = (int)t.shape[0], n_out = t.shape.size() > 1 ? (int)t.shape[1] : 1; + DqArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)ref::row_bytes(t.type, n_in), perm}; + D(P(lib, "dq_f16_" + std::to_string(t.type)), {mtl::Arg::B(W, t.offset), mtl::Arg::B(out, out_off), mtl::Arg::V(a)}, + (n_in / unit_of(t.type) + 63) / 64, n_out, 64, "dq"); + }; + auto act16 = [&](const mtl::Buffer& x, const mtl::Buffer& xh, size_t n, const std::string& c) { + const uint32_t nn = (uint32_t)n; + D(P(lib, "act_f16"), {mtl::Arg::B(x), mtl::Arg::B(xh), mtl::Arg::V(nn)}, (nn + 255) / 256, 1, 256, c); + }; + int cur_inj = 0; + // the HC read of T rows: mixed (fp32) and mixh (fp16) + auto hc_read = [&](const std::string& pre, bool apply, int T) { + const gguf::Tensor &tn = T_(pre + "_norm.weight"), &tdn = T_(pre + "_down.weight"), &tin = T_(pre + "_inject.weight"), + &tup = T_(pre + "_up.weight"); + dq(tdn, WD); + dq(tin, WD, (size_t)2 * LR * HC * N); + GrArgs a{N, HC, LR, (uint32_t)T, apply ? 1u : 0u, 1e-6f}; + EpiArgs ea{(uint32_t)T, N, HC, LR}; + if (hc_mode == "fused") { + // fp16 xn only; the up GEMM's epilogue mixes (w_up rows interleaved by the dequantization) + dq(tup, WU, 0, HC); + D(P(lib, "gr_norm_h"), {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj[cur_inj]), mtl::Arg::B(W, tn.offset), + mtl::Arg::B(xnh), mtl::Arg::V(a)}, HC, T, 256, "hc"); + gemm(xnh, WD, V, T, LR + HC, HC * N, "hc gemm"); + D(P(lib, "gr_epi_down"), {mtl::Arg::B(V), mtl::Arg::B(loh), mtl::Arg::B(inj[1 - cur_inj]), mtl::Arg::V(ea)}, + (T * (LR + HC) + 255) / 256, 1, 256, "hc"); + GemmArgs g{T, HC * N, LR, 0}; + D(P(lib, "gemm_hc_up_64_64_4"), {mtl::Arg::B(loh), mtl::Arg::B(WU), mtl::Arg::B(xnh), mtl::Arg::B(mixed), + mtl::Arg::B(mixh), mtl::Arg::V(g), mtl::Arg::V(ea)}, + HC * N / 64, (T + 63) / 64, 128, "hc gemm"); + cur_inj = 1 - cur_inj; + return; + } + dq(tup, WU); + D(P(lib, "gr_norm"), {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj[cur_inj]), mtl::Arg::B(W, tn.offset), + mtl::Arg::B(xn), mtl::Arg::V(a)}, HC, T, 256, "hc"); + act16(xn, xnh, (size_t)T * HC * N, "hc"); + gemm(xnh, WD, V, T, LR + HC, HC * N, "hc gemm"); + D(P(lib, "gr_epi_down"), {mtl::Arg::B(V), mtl::Arg::B(loh), mtl::Arg::B(inj[1 - cur_inj]), mtl::Arg::V(ea)}, + (T * (LR + HC) + 255) / 256, 1, 256, "hc"); + gemm(loh, WU, G, T, HC * N, LR, "hc gemm"); + D(P(lib, "gr_epi_up"), {mtl::Arg::B(G), mtl::Arg::B(xn), mtl::Arg::B(mixed), mtl::Arg::B(mixh), mtl::Arg::V(ea)}, + (T * N + 255) / 256, 1, 256, "hc"); + cur_inj = 1 - cur_inj; + }; + auto build = [&](int pos0, int T, bool head) { + ds.clear(); + cls.clear(); + int gi = 0, qi = 0; + for (int l = 0; l < 48; ++l) { + const std::string b = "blk." + std::to_string(l) + "."; + hc_read(b + "hc_attn", l > 0, T); + if (l % 4 != 3) { + const gguf::Tensor &tq = T_(b + "attn_qkv.weight"), &tz = T_(b + "attn_gate.weight"), + &ta = T_(b + "ssm_alpha.weight"), &tb = T_(b + "ssm_beta.weight"), &to = T_(b + "ssm_out.weight"); + dq(tq, Wqkv); + dq(tz, Wz); + dq(ta, Wal); + dq(tb, Wbe); + dq(to, Wout); + gemm(mixh, Wqkv, qkv, T, 10240, N, "gdn proj"); + gemm(mixh, Wz, z, T, 6144, N, "gdn proj"); + gemm(mixh, Wal, al, T, 48, N, "gdn proj"); + gemm(mixh, Wbe, be, T, 48, N, "gdn proj"); + GdnArgs a{16, 48, 10240, 1, 1e-6f, 0}; + const int par = gpar[gi]; + if (gdn_mode != "step") { + GdnPArgs ga{(uint32_t)T, 16, 48, 10240, 1e-6f, 0}; + D(P(lgdn, "gdn_pre"), + {mtl::Arg::B(qkv), mtl::Arg::B(ghist[par][gi]), mtl::Arg::B(W, T_(b + "ssm_conv1d.weight").offset), + mtl::Arg::B(al), mtl::Arg::B(be), mtl::Arg::B(W, T_(b + "ssm_dt.bias").offset), + mtl::Arg::B(W, T_(b + "ssm_a").offset), mtl::Arg::B(gqn), mtl::Arg::B(gkn), mtl::Arg::B(gvc), mtl::Arg::B(ggb), + mtl::Arg::B(gkq), mtl::Arg::B(ghist[1 - par][gi]), mtl::Arg::V(ga)}, + 16 + 24, T, 256, "gdn scan"); + const int ib = gdn_mode == "scan_v" ? 8 : std::atoi(gdn_mode.c_str() + 6); + D(P(lgdn, "gdn_" + gdn_mode, gdn_mode == "scan_v" ? mtl::Constants() : mtl::Constants().i(6, 0)), + {mtl::Arg::B(gst[gi]), mtl::Arg::B(gqn), mtl::Arg::B(gkn), mtl::Arg::B(gvc), mtl::Arg::B(ggb), mtl::Arg::B(gkq), + mtl::Arg::B(go), mtl::Arg::V(ga)}, + 48 * 4, 1, 32 * ib, "gdn scan"); + D(P(lgdn, "gdn_post"), {mtl::Arg::B(go), mtl::Arg::B(z), mtl::Arg::B(W, T_(b + "ssm_norm.weight").offset), + mtl::Arg::B(ygh), mtl::Arg::B(yg), mtl::Arg::V(ga)}, + 48, T, 128, "gdn scan"); + gemm(ygh, Wout, bo, T, N, 6144, "gdn proj"); + ++gi; + goto moe; + } + D(P(lgdn, "gdn_step", mtl::Constants().i(2, T), "gdn_step#" + std::to_string(T)), + {mtl::Arg::B(gst[gi]), mtl::Arg::B(ghist[par][gi]), mtl::Arg::B(qkv), mtl::Arg::B(al), mtl::Arg::B(be), + mtl::Arg::B(z), mtl::Arg::B(W, T_(b + "ssm_conv1d.weight").offset), mtl::Arg::B(W, T_(b + "ssm_dt.bias").offset), + mtl::Arg::B(W, T_(b + "ssm_a").offset), mtl::Arg::B(W, T_(b + "ssm_norm.weight").offset), mtl::Arg::B(yg), + mtl::Arg::V(a), mtl::Arg::B(ghist[1 - par][gi])}, + 48, 1, 512, "gdn scan"); + act16(yg, ygh, (size_t)T * 6144, "gdn proj"); + gemm(ygh, Wout, bo, T, N, 6144, "gdn proj"); + ++gi; + } else { + const gguf::Tensor &tq = T_(b + "attn_q.weight"), &tk = T_(b + "attn_k.weight"), &tv = T_(b + "attn_v.weight"), + &tiq = T_(b + "indexer.q_proj.weight"), &tik = T_(b + "indexer.k_proj.weight"), + &to = T_(b + "attn_output.weight"); + dq(tq, Wq); + dq(tk, Wk); + dq(tv, Wv); + dq(tiq, Wiq); + dq(tik, Wik); + dq(to, Wout); + gemm(mixh, Wq, qfull, T, 12288, N, "qsa proj"); + gemm(mixh, Wk, kcur, T, 512, N, "qsa proj"); + gemm(mixh, Wv, vcur, T, 512, N, "qsa proj"); + gemm(mixh, Wiq, iqr, T, 512, N, "qsa proj"); + gemm(mixh, Wik, ikr, T, 128, N, "qsa proj"); + QsaArgs a{(uint32_t)pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f}; + D(P(lqsa, "qsa_prep"), + {mtl::Arg::B(qfull), mtl::Arg::B(kcur), mtl::Arg::B(vcur), mtl::Arg::B(iqr), mtl::Arg::B(ikr), + mtl::Arg::B(W, T_(b + "attn_q_norm.weight").offset), mtl::Arg::B(W, T_(b + "attn_k_norm.weight").offset), + mtl::Arg::B(W, T_(b + "indexer.q_norm.weight").offset), mtl::Arg::B(W, T_(b + "indexer.k_norm.weight").offset), + mtl::Arg::B(cs), mtl::Arg::B(qcur), mtl::Arg::B(qidx), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]), + mtl::Arg::B(ksc[qi]), mtl::Arg::B(vsc[qi]), mtl::Arg::B(tail[qi]), mtl::Arg::B(dead[qi]), + mtl::Arg::B(pooled[qi]), mtl::Arg::V(a)}, + 33, T, 256, "qsa prep"); + if (pos0 + T > QB) { // some token's context is past the budget: select + if (!std::getenv("PC_QSA_SCORE")) // the GEMM score unless PC_QSA_SCORE (= the scalar qsa_score) + D(P(lqsa, "qsa_score_mm"), {mtl::Arg::B(qidx), mtl::Arg::B(pooled[qi]), mtl::Arg::B(dead[qi]), + mtl::Arg::B(scores), mtl::Arg::V(a)}, + ((pos0 + T) / 4 + 1 + 63) / 64, (T + 15) / 16, 128, "qsa score"); + else + D(P(lqsa, "qsa_score"), {mtl::Arg::B(qidx), mtl::Arg::B(pooled[qi]), mtl::Arg::B(dead[qi]), mtl::Arg::B(scores), + mtl::Arg::V(a)}, + ((pos0 + T) / 4 + 1 + 7) / 8, T, 256, "qsa score"); + D(P(lqsa, "qsa_topk"), {mtl::Arg::B(scores), mtl::Arg::B(qids), mtl::Arg::V(a)}, T, 1, 1024, "qsa sel"); + } else { + // within the budget the selection is the identity: qsa_topk only writes 0..n_kv-1 + D(P(lqsa, "qsa_topk"), {mtl::Arg::B(scores), mtl::Arg::B(qids), mtl::Arg::V(a)}, T, 1, 1024, "qsa sel"); + } + if (qsa_mode == "flash") { + const uint32_t nq = (uint32_t)T * 24 * 256; + D(P(lqsa, "qsa_q_f16"), {mtl::Arg::B(qcur), mtl::Arg::B(qcurh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 256, "qsa attn"); + D(P(lqsa, "qsa_flash", mtl::Constants().b(7, k4)), + {mtl::Arg::B(qcurh), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]), mtl::Arg::B(ksc[qi]), mtl::Arg::B(vsc[qi]), + mtl::Arg::B(qids), mtl::Arg::B(qfull), mtl::Arg::B(xqh), mtl::Arg::V(a)}, + 2, T, 128, "qsa attn"); + gemm(xqh, Wout, bo, T, N, 6144, "qsa proj"); + ++qi; + goto moe; + } + if (qsa_mode == "mm") { + const uint32_t nq = (uint32_t)T * 24 * 256; + D(P(lqsa, "qsa_q_f16"), {mtl::Arg::B(qcur), mtl::Arg::B(qcurh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 256, "qsa attn"); + D(P(lqsa, "qsa_attn_mm", mtl::Constants().b(7, k4)), {mtl::Arg::B(qcurh), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]), mtl::Arg::B(ksc[qi]), + mtl::Arg::B(vsc[qi]), mtl::Arg::B(qids), mtl::Arg::B(part), mtl::Arg::V(a)}, + QCHUNKS, 2, 32, "qsa attn", T); + } else + D(P(lqsa, "qsa_attn"), {mtl::Arg::B(qcur), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]), mtl::Arg::B(ksc[qi]), + mtl::Arg::B(vsc[qi]), mtl::Arg::B(qids), mtl::Arg::B(part), mtl::Arg::V(a)}, + QCHUNKS, 2, 256, "qsa attn", T); + D(P(lqsa, "qsa_merge"), {mtl::Arg::B(part), mtl::Arg::B(qfull), mtl::Arg::B(XQ), mtl::Arg::V(a)}, 24, T, 256, + "qsa merge"); + act16(XQ, xqh, (size_t)T * 6144, "qsa proj"); + gemm(xqh, Wout, bo, T, N, 6144, "qsa proj"); + ++qi; + } + moe: + // ---- MoE + hc_read(b + "hc_ffn", true, T); + const gguf::Tensor &tr = T_(b + "ffn_gate_inp.weight"), &tsg = T_(b + "ffn_gate_inp_shexp.weight"), + &tg_ = T_(b + "ffn_gate_shexp.weight"), &tu_ = T_(b + "ffn_up_shexp.weight"), + &td_ = T_(b + "ffn_down_shexp.weight"), &eg = T_(b + "ffn_gate_exps.weight"), + &eu = T_(b + "ffn_up_exps.weight"), &ed = T_(b + "ffn_down_exps.weight"); + dq(tr, Wr); + dq(tsg, Wsg); + dq(tg_, Wshg); + dq(tu_, Wshu); + dq(td_, Wshd); + gemm(mixh, Wr, logits, T, NEXP, N, "router"); + const uint32_t nt = T; + D(P(lib, "route_topk"), {mtl::Arg::B(logits), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1, + 256, "router"); + const int nent = T * TOPK, nblk = (nent + 1023) / 1024; + MmPlanArgs pa{(uint32_t)nent, TOPK, NEXP, (uint32_t)TM, (uint32_t)nblk}; + const uint32_t nclr = maxrows; + D(P(lib, "mm_clear"), {mtl::Arg::B(stok), mtl::Arg::V(nclr)}, (maxrows + 255) / 256, 1, 256, "experts"); + D(P(lib, "mm_hist"), {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024, "experts"); + D(P(lib, "mm_scan"), {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), + mtl::Arg::V(pa)}, 1, 1, NEXP, "experts"); + D(P(lib, "mm_place"), {mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), + mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}, nblk, 1, 1024, "experts"); + MmArgs ma{eg.offset, eu.offset, ed.offset, 0, 0, (uint32_t)ref::row_bytes(eg.type, N), + (uint32_t)ref::row_bytes(ed.type, NFF), N, NFF}; + ma.gu_stride = (uint64_t)ma.gu_row * NFF; + ma.d_stride = (uint64_t)ma.d_row * N; + const uint32_t ne = N; + mtl::Dispatch g; + g.p = P(lib, "mm_gather"); + g.args = {mtl::Arg::B(mixed), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; + g.tg[0] = 256; + g.indirect = &meta; + g.ind_off = 8 * 4; + push(g, "experts"); + mtl::Dispatch gu; + gu.p = gu_v3 ? P(lib, "mm_gu3_" + std::to_string(eg.type) + "_" + gu3_for(eg.type), mtl::Constants().b(7, k4)) + : P(lib, "mm_gu2_" + std::to_string(eg.type) + (fmt_unit32(eg.type) ? "_16_64_32_2" : "_16_64_16_2"), + mtl::Constants().i(5, 0)); + gu.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + gu.tg[0] = gu_v3 ? 32 * gu3[3] : 64; + gu.indirect = &meta; + gu.ind_off = 2 * 4; + push(gu, "experts gu"); + mtl::Dispatch dn; + dn.p = dn_v3 ? P(lib, "mm_down3_" + std::to_string(ed.type) + "_" + tl3(dn3), mtl::Constants().b(7, k4)) + : P(lib, "mm_down2_" + std::to_string(ed.type) + (ed.type == ref::IQ4_NL ? "_16_64_32_2" : "_16_64_16_2")); + dn.args = {mtl::Arg::B(W), mtl::Arg::B(HS), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)}; + dn.tg[0] = dn_v3 ? 32 * dn3[3] : 64; + dn.indirect = &meta; + dn.ind_off = 5 * 4; + push(dn, "experts dn"); + const uint32_t ca[4] = {(uint32_t)T, N, TOPK, 0}; + D(P(lib, "mm_combine"), {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(yr), mtl::Arg::V(ca)}, + (T * N + 255) / 256, 1, 256, "experts"); + gemm(mixh, Wshg, gsh, T, NFF, N, "shared"); + gemm(mixh, Wshu, ush, T, NFF, N, "shared"); + const uint32_t nh = T * NFF; + D(P(lib, "swiglu_h"), {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(hh), mtl::Arg::V(nh)}, (nh + 255) / 256, 1, 256, + "shared"); + gemm(hh, Wshd, ys, T, N, NFF, "shared"); + gemm(mixh, Wsg, sgv, T, 1, N, "shared"); + const uint32_t oa[2] = {(uint32_t)T, N}; + D(P(lib, "moe_out"), {mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sgv), mtl::Arg::B(bo), mtl::Arg::V(oa)}, + (T * N + 255) / 256, 1, 256, "shared"); + } + if (head) { + // the final read of the last row (decode kernels, 1 token) and the LM head + const size_t lr = (size_t)T - 1; + const gguf::Tensor &tn = T_("output_hc_norm.weight"), &tdn = T_("output_hc_down.weight"), &tup = T_("output_hc_up.weight"); + GrArgs a{N, HC, LR, 1, 1, 1e-6f}; + D(P(lib, "gr_norm"), {mtl::Arg::B(R, lr * HC * N * 4), mtl::Arg::B(bo, lr * N * 4), mtl::Arg::B(inj[cur_inj], lr * HC * 4), + mtl::Arg::B(W, tn.offset), mtl::Arg::B(xn1), mtl::Arg::V(a)}, HC, 1, 256, "head"); + D(P(lib, "gr_down_8_8", mtl::Constants().i(1, 1), "gr_down_8_8#1"), + {mtl::Arg::B(W, tdn.offset), mtl::Arg::B(W, tdn.offset), mtl::Arg::B(xn1), mtl::Arg::B(lo1), mtl::Arg::B(inj1), + mtl::Arg::V(a)}, LR + HC, 1, 256, "head"); + D(P(lib, "gr_up", mtl::Constants().i(1, 1), "gr_up#1"), + {mtl::Arg::B(W, tup.offset), mtl::Arg::B(lo1), mtl::Arg::B(xn1), mtl::Arg::B(X1), mtl::Arg::V(a)}, N / 16, 1, 256, "head"); + const gguf::Tensor& to = T_("output.weight"); + GemvArgs ga{N, NV, 1, (uint32_t)ref::row_bytes(to.type, N), N, NV}; + const std::string fn = "gemv_" + std::to_string(to.type) + "_16_1_4_1"; + const int xk = to.type == ref::Q4_K || to.type == ref::Q5_K ? 2 : (to.type == ref::Q6_K || to.type == ref::Q3_K) ? 3 : 0; + D(P(lib, fn, mtl::Constants().i(0, 1), fn + "#1"), + {mtl::Arg::B(W, to.offset), mtl::Arg::B(X1, (size_t)4 * N * xk), mtl::Arg::B(X1, (size_t)4 * N * 4), mtl::Arg::B(lm), + mtl::Arg::V(ga)}, (NV + 7) / 8, 1, 128, "head"); + ArgArgs aa{NV, 120}; + D(P(lsmp, "argmax_part"), {mtl::Arg::B(lm), mtl::Arg::B(apv), mtl::Arg::B(api), mtl::Arg::V(aa)}, 120, 1, 256, "head"); + D(P(lsmp, "argmax_final"), {mtl::Arg::B(apv), mtl::Arg::B(api), mtl::Arg::B(tokout), mtl::Arg::V(aa)}, 1, 1, 256, "head"); + } + }; + // a chunk: the embedding rows into every stream, then the chunk's command buffer; GDN history ping-pong flipped + auto run_chunk = [&](int pos0, int T, bool head) { + for (int t = 0; t < T; ++t) + for (int c = 0; c < HC; ++c) std::memcpy(R.as() + ((size_t)t * HC + c) * N, E.data() + (size_t)t * N, 4 * N); + cur_inj = 0; + build(pos0, T, head); + const double g = cx.run(ds); + for (auto& p : gpar) p ^= 1; + return g; + }; + auto prompt = [&](int Pn, double* first_chunk) { + double tot = 0; + for (int p0 = 0; p0 < Pn; p0 += C) { + const int T = std::min(C, Pn - p0); + const double g = run_chunk(p0, T, p0 + T >= Pn); + if (p0 == 0 && first_chunk) *first_chunk = g; + tot += g; + } + return tot; + }; + std::set all_cls; + // warm: the first chunk until the clocks are up + { + const double t0 = mtl::now_s(); + int n = 0; + while (mtl::now_s() - t0 < 1.5) { run_chunk(0, std::min(C, prompts[0]), false); ++n; } + std::printf("warm-up: %d chunks; %zu dispatches a chunk\n", n, ds.size()); + all_cls.insert(cls.begin(), cls.end()); + } + if (breakdown) { + // each class alone on the first chunk (pos0 0) and on a chunk at pos0 = PMAX - C (QSA past the budget) + if (PMAX > C) { // the classes of a chunk past the budget (score) too + run_chunk(std::max(0, PMAX - C), C, false); + all_cls.insert(cls.begin(), cls.end()); + } + for (int p0 : {0, std::max(0, PMAX - C)}) { + const std::string keep_only = only; + std::map by; + double all = 0; + std::vector names = {""}; + for (const auto& c : all_cls) names.push_back(c); + for (const std::string& c : names) { + only = c; + std::vector t; + for (int i = 0; i < 3; ++i) t.push_back(run_chunk(p0, C, false)); + if (c.empty()) all = median(t); + else by[c] = median(t); + } + only = keep_only; + double sum = 0; + for (auto& kv : by) sum += kv.second; + std::printf("chunk of %d at pos %d: all %.2f ms (%.0f tok/s); classes alone (sum %.2f ms):\n", C, p0, all * 1e3, C / all, + sum * 1e3); + for (auto& kv : by) std::printf(" %-10s %8.2f ms %6.2f us a token\n", kv.first.c_str(), kv.second * 1e3, kv.second / C * 1e6); + } + } + for (int Pn : prompts) { + std::vector t; + double fc = 0; + for (int r = 0; r < reps; ++r) t.push_back(prompt(Pn, &fc)); + const double m = median(t); + std::printf("prompt %6d tokens (chunks of %d): GPU %8.1f ms -> %7.1f tok/s (first chunk %.1f ms)\n", Pn, C, m * 1e3, + Pn / m, fc * 1e3); + } + return 0; +} diff --git a/edge0-strata/metal/tests/prompt_parity.cpp b/edge0-strata/metal/tests/prompt_parity.cpp new file mode 100644 index 00000000..2cb48351 --- /dev/null +++ b/edge0-strata/metal/tests/prompt_parity.cpp @@ -0,0 +1,222 @@ +// metal/tests/prompt_parity.cpp - the prompt-chunk glue (prompt.metal) against CPU references: +// route_topk softmax + top 10 + renormalised weights (moe_route's rule) for 300 tokens, a third of them +// with exact ties +// the HC read as GEMMs gr_norm -> act_f16 -> [w_down; w_inject] GEMM -> gr_epi_down -> w_up GEMM -> gr_epi_up, on +// random BF16 weights, vs the exact formula (fp64) - the fp16 operands set the tolerance +// swiglu_h, moe_out elementwise +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include + +namespace { +constexpr int N = 2560, HC = 4, LR = 320; +struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; +struct EpiArgs { uint32_t n_tok, n_embd, hc, hc_lr; }; +struct DqArgs { uint32_t n_in, n_out, row_bytes, perm; }; +struct GemmArgs { int M, N, K, pad; }; +double sig(double v) { return 1.0 / (1.0 + std::exp(-v)); } +} // namespace + +int main() { + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "gr.metal", "prompt.metal"}); + std::printf("device %s\n", cx.device_name().c_str()); + int fails = 0; + // ---- route_topk + { + const int T = 300; + mtl::Pipeline p = cx.pipeline(lib, "route_topk"); + mtl::Buffer lg = cx.alloc(4 * T * 512), ids = cx.alloc(4 * T * 10), w = cx.alloc(4 * T * 10); + ref::Rng r(9); + for (int t = 0; t < T; ++t) + for (int e = 0; e < 512; ++e) { + float v = r.sym() * 4.0f; + if (t % 3 == 0) v = std::round(v * 4.0f) / 4.0f; + lg.as()[t * 512 + e] = v; + } + const uint32_t nt = T; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::V(nt)}; + d.grid[0] = (T + 7) / 8; + d.tg[0] = 256; + cx.run({d}); + int bad = 0; + double worst = 0; + for (int t = 0; t < T; ++t) { + const float* l = lg.as() + t * 512; + float mx = -INFINITY; + for (int e = 0; e < 512; ++e) mx = std::max(mx, l[e]); + std::vector pr(512); + double s = 0; + for (int e = 0; e < 512; ++e) s += (pr[e] = std::exp((double)l[e] - mx)); + std::vector order(512); + for (int e = 0; e < 512; ++e) order[e] = e; + std::stable_sort(order.begin(), order.end(), [&](int x, int y) { return l[x] > l[y]; }); + double ss = 0; + for (int k = 0; k < 10; ++k) ss += pr[order[k]] / s; + for (int k = 0; k < 10; ++k) { + if (ids.as()[t * 10 + k] != order[k]) { ++bad; break; } + const double ref = pr[order[k]] / s / std::max(ss, 6.103515625e-5); + worst = std::max(worst, std::fabs(w.as()[t * 10 + k] - ref) / ref); + } + } + const bool ok = bad == 0 && worst < 2e-6; + fails += !ok; + std::printf("route_topk, %d tokens: id mismatches %d, worst weight rel err %.2e %s\n", T, bad, worst, ok ? "PASS" : "FAIL"); + } + // ---- the HC read as GEMMs + for (int apply = 0; apply < 2; ++apply) { + const int T = 70, NIN = HC * N, NR = LR + HC; + mtl::Buffer R = cx.alloc(4 * (size_t)T * NIN), bo = cx.alloc(4 * (size_t)T * N), inj = cx.alloc(4 * T * HC), + wn = cx.alloc(4 * NIN), wdn = cx.alloc(2 * (size_t)LR * NIN), win = cx.alloc(2 * (size_t)HC * NIN), + wup = cx.alloc(2 * (size_t)NIN * LR), xn = cx.alloc(4 * (size_t)T * NIN), xnh = cx.alloc(2 * (size_t)T * NIN), + WD = cx.alloc(2 * (size_t)NR * NIN), WU = cx.alloc(2 * (size_t)NIN * LR), V = cx.alloc(4 * (size_t)T * NR), + loh = cx.alloc(2 * (size_t)T * LR), injo = cx.alloc(4 * T * HC), G = cx.alloc(4 * (size_t)T * NIN), + mixed = cx.alloc(4 * (size_t)T * N), mixh = cx.alloc(2 * (size_t)T * N); + ref::Rng r(11 + apply); + std::vector R0((size_t)T * NIN); + for (auto& v : R0) v = r.sym(); + std::copy(R0.begin(), R0.end(), R.as()); + for (size_t i = 0; i < (size_t)T * N; ++i) bo.as()[i] = r.sym(); + for (int i = 0; i < T * HC; ++i) inj.as()[i] = r.sym() * 3.0f; + for (int i = 0; i < NIN; ++i) wn.as()[i] = 0.5f + r.uni(); + for (size_t i = 0; i < (size_t)LR * NIN; ++i) wdn.as()[i] = ref::f32_to_bf16(r.sym() * 0.05f); + for (size_t i = 0; i < (size_t)HC * NIN; ++i) win.as()[i] = ref::f32_to_bf16(r.sym() * 0.05f); + for (size_t i = 0; i < (size_t)NIN * LR; ++i) wup.as()[i] = ref::f32_to_bf16(r.sym() * 0.2f); + mtl::Pipeline pn = cx.pipeline(lib, "gr_norm"), pa = cx.pipeline(lib, "act_f16"), pdq = cx.pipeline(lib, "dq_f16_30"), + pg = cx.pipeline(lib, "gemm_f16_64_64_4"), ped = cx.pipeline(lib, "gr_epi_down"), peu = cx.pipeline(lib, "gr_epi_up"); + std::vector ds; + GrArgs ga{N, HC, LR, (uint32_t)T, (uint32_t)apply, 1e-6f}; + EpiArgs ea{(uint32_t)T, N, HC, LR}; + auto D = [&](const mtl::Pipeline& p, std::vector a, uint32_t gx, uint32_t gy, uint32_t tg) { + mtl::Dispatch d; d.p = &p; d.args = std::move(a); d.grid[0] = gx; d.grid[1] = gy; d.tg[0] = tg; ds.push_back(d); + }; + auto dq = [&](const mtl::Buffer& w, const mtl::Buffer& out, size_t ooff, int n_in, int n_out) { + DqArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)(2 * n_in)}; + D(pdq, {mtl::Arg::B(w), mtl::Arg::B(out, ooff), mtl::Arg::V(a)}, (n_in / 16 + 63) / 64, n_out, 64); + }; + auto gemm = [&](const mtl::Buffer& A, const mtl::Buffer& B, const mtl::Buffer& Y, int M, int Nn, int K) { + GemmArgs g{M, Nn, K, 0}; + D(pg, {mtl::Arg::B(A), mtl::Arg::B(B), mtl::Arg::B(Y), mtl::Arg::V(g)}, (Nn + 63) / 64, (M + 63) / 64, 128); + }; + const uint32_t nx = T * NIN; + dq(wdn, WD, 0, NIN, LR); + dq(win, WD, (size_t)2 * LR * NIN, NIN, HC); + dq(wup, WU, 0, LR, NIN); + D(pn, {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xn), mtl::Arg::V(ga)}, HC, T, 256); + D(pa, {mtl::Arg::B(xn), mtl::Arg::B(xnh), mtl::Arg::V(nx)}, (nx + 255) / 256, 1, 256); + gemm(xnh, WD, V, T, NR, NIN); + D(ped, {mtl::Arg::B(V), mtl::Arg::B(loh), mtl::Arg::B(injo), mtl::Arg::V(ea)}, (T * NR + 255) / 256, 1, 256); + gemm(loh, WU, G, T, NIN, LR); + D(peu, {mtl::Arg::B(G), mtl::Arg::B(xn), mtl::Arg::B(mixed), mtl::Arg::B(mixh), mtl::Arg::V(ea)}, (T * N + 255) / 256, 1, 256); + cx.run(ds); + // the fused read: gr_norm_h (fp16 xn), the down GEMM, gr_epi_down, the up GEMM with the epilogue (w_up rows + // interleaved by dq_f16's perm) + mtl::Buffer mixed2 = cx.alloc(4 * (size_t)T * N), mixh2 = cx.alloc(2 * (size_t)T * N), WU2 = cx.alloc(2 * (size_t)NIN * LR), + xnh2 = cx.alloc(2 * (size_t)T * NIN), V2 = cx.alloc(4 * (size_t)T * NR), loh2 = cx.alloc(2 * (size_t)T * LR), + injo2 = cx.alloc(4 * T * HC); + std::copy(R0.begin(), R0.end(), R.as()); + { + mtl::Pipeline pnh = cx.pipeline(lib, "gr_norm_h"), pup = cx.pipeline(lib, "gemm_hc_up_64_64_4"); + ds.clear(); + DqArgs da{(uint32_t)LR, (uint32_t)NIN, (uint32_t)(2 * LR), (uint32_t)HC}; + D(pdq, {mtl::Arg::B(wup), mtl::Arg::B(WU2), mtl::Arg::V(da)}, (LR / 16 + 63) / 64, NIN, 64); + D(pnh, {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj), mtl::Arg::B(wn), mtl::Arg::B(xnh2), mtl::Arg::V(ga)}, HC, T, 256); + gemm(xnh2, WD, V2, T, NR, NIN); + D(ped, {mtl::Arg::B(V2), mtl::Arg::B(loh2), mtl::Arg::B(injo2), mtl::Arg::V(ea)}, (T * NR + 255) / 256, 1, 256); + GemmArgs g{T, NIN, LR, 0}; + D(pup, {mtl::Arg::B(loh2), mtl::Arg::B(WU2), mtl::Arg::B(xnh2), mtl::Arg::B(mixed2), mtl::Arg::B(mixh2), mtl::Arg::V(g), + mtl::Arg::V(ea)}, NIN / 64, (T + 63) / 64, 128); + cx.run(ds); + } + // reference + double se = 0, sr = 0, sei = 0, sri = 0, ser = 0, srr = 0, se2 = 0, sh2 = 0; + std::vector x(NIN), lo(LR); + for (int t = 0; t < T; ++t) { + for (int c = 0; c < HC; ++c) { + const double g = apply ? 2.0 * sig(inj.as()[t * HC + c] / (double)HC) : 0.0; + double ss = 0; + for (int d = 0; d < N; ++d) { + const size_t i = (size_t)t * NIN + c * N + d; + const double v = R0[i] + (apply ? bo.as()[(size_t)t * N + d] * g : 0.0); + x[c * N + d] = v; + ss += v * v; + ser += (R.as()[i] - v) * (R.as()[i] - v); + srr += v * v; + } + const double rs = 1.0 / std::sqrt(ss / N + 1e-6); + for (int d = 0; d < N; ++d) x[c * N + d] *= wn.as()[c * N + d] * rs; + } + for (int k = 0; k < NR; ++k) { + const uint16_t* w = k < LR ? wdn.as() + (size_t)k * NIN : win.as() + (size_t)(k - LR) * NIN; + double s = 0; + for (int i = 0; i < NIN; ++i) s += ref::bf16_to_f32(w[i]) * x[i]; + if (k < LR) lo[k] = s / HC * sig(s / HC); + else { + const double gv = injo.as()[t * HC + k - LR]; + sei += (gv - s) * (gv - s); + sri += s * s; + } + } + for (int d = 0; d < N; ++d) { + double m = 0; + for (int c = 0; c < HC; ++c) { + const uint16_t* w = wup.as() + (size_t)(c * N + d) * LR; + double s = 0; + for (int k = 0; k < LR; ++k) s += ref::bf16_to_f32(w[k]) * lo[k]; + m += x[c * N + d] * sig(s); + } + m /= HC; + const double gv = mixed.as()[(size_t)t * N + d], g2 = mixed2.as()[(size_t)t * N + d], + h2 = ref::f16_to_f32(mixh2.as()[(size_t)t * N + d]); + se += (gv - m) * (gv - m); + se2 += (g2 - m) * (g2 - m); + sh2 += (h2 - m) * (h2 - m); + sr += m * m; + } + } + const double rm = std::sqrt(se / sr), ri = std::sqrt(sei / sri), rr = std::sqrt(ser / srr); + const double r2 = std::sqrt(se2 / sr), rh2 = std::sqrt(sh2 / sr); + const bool ok = rm < 2e-3 && ri < 2e-3 && rr < 1e-6 && r2 < 2e-3 && rh2 < 2e-3; + fails += !ok; + std::printf("HC read as GEMMs, %d tokens, apply %d: mixed rel err %.2e, inject %.2e, residual %.2e; fused read " + "(fp16 xn, up GEMM epilogue): mixed %.2e, fp16 %.2e %s\n", T, apply, rm, ri, rr, r2, rh2, ok ? "PASS" : "FAIL"); + } + // ---- swiglu_h, moe_out + { + const int n = 5000, T = 7, NE = 640; + mtl::Pipeline ps = cx.pipeline(lib, "swiglu_h"), pm = cx.pipeline(lib, "moe_out"); + mtl::Buffer g = cx.alloc(4 * n), u = cx.alloc(4 * n), h = cx.alloc(2 * n), yr = cx.alloc(4 * T * NE), + ys = cx.alloc(4 * T * NE), sg = cx.alloc(4 * T), bo = cx.alloc(4 * T * NE); + ref::Rng r(5); + for (int i = 0; i < n; ++i) { g.as()[i] = r.sym() * 6; u.as()[i] = r.sym(); } + for (int i = 0; i < T * NE; ++i) { yr.as()[i] = r.sym(); ys.as()[i] = r.sym(); } + for (int i = 0; i < T; ++i) sg.as()[i] = r.sym() * 4; + const uint32_t nn = n; + const uint32_t a2[2] = {(uint32_t)T, (uint32_t)NE}; + mtl::Dispatch d1; d1.p = &ps; d1.args = {mtl::Arg::B(g), mtl::Arg::B(u), mtl::Arg::B(h), mtl::Arg::V(nn)}; + d1.grid[0] = (n + 255) / 256; d1.tg[0] = 256; + mtl::Dispatch d2; d2.p = ± d2.args = {mtl::Arg::B(yr), mtl::Arg::B(ys), mtl::Arg::B(sg), mtl::Arg::B(bo), mtl::Arg::V(a2)}; + d2.grid[0] = (T * NE + 255) / 256; d2.tg[0] = 256; + cx.run({d1, d2}); + double w1 = 0, w2 = 0; + for (int i = 0; i < n; ++i) { + const double gv = g.as()[i], ref = gv * sig(gv) * u.as()[i]; + w1 = std::max(w1, std::fabs(ref::f16_to_f32(h.as()[i]) - ref) / (std::fabs(ref) + 1e-3)); + } + for (int i = 0; i < T * NE; ++i) { + const double ref = yr.as()[i] + sig(sg.as()[i / NE]) * ys.as()[i]; + w2 = std::max(w2, std::fabs(bo.as()[i] - ref)); + } + const bool ok = w1 < 1e-3 && w2 < 1e-5; + fails += !ok; + std::printf("swiglu_h worst rel err %.2e, moe_out worst abs err %.2e %s\n", w1, w2, ok ? "PASS" : "FAIL"); + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/qsa_parity.cpp b/edge0-strata/metal/tests/qsa_parity.cpp new file mode 100644 index 00000000..8162ec2e --- /dev/null +++ b/edge0-strata/metal/tests/qsa_parity.cpp @@ -0,0 +1,308 @@ +// metal/tests/qsa_parity.cpp - QSA decode (qsa.metal) stage by stage against a CPU reference, and its speed. +// +// qsa_parity [P] a context of P cells built by qsa_prep (default 3000: the selection is active above +// 2051), then a window of 3 tokens: prep, score, top-k, attention, merge. Each stage is +// checked against the CPU fed with the GPU's inputs to that stage; the ids exactly. +// qsa_parity --bench [P] GPU time of the window's kernels at context P, 1 and 3 tokens +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int QH = 24, QKV = 2, QD = 256, IH = 4, ID = 128, BUDGET = 2051, CH = 64, NCH = (BUDGET + CH - 1) / CH; +struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +struct Err { + double se = 0, sr = 0; + void add(double g, double r) { se += (g - r) * (g - r); sr += r * r; } + double l2() const { return sr > 0 ? std::sqrt(se / sr) : std::sqrt(se); } +}; +float h2f(uint16_t h) { return ref::f16_to_f32(h); } +} // namespace + +int main(int argc, char** argv) { + bool bench = false; + int P = 3000; + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--bench") bench = true; + else P = std::atoi(argv[i]); + } + const int MAXT = 8, MAXC = std::max(4096, (P + 64) / 4 * 4 + 64); + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa2.metal"}, "#define NO_GGML_COMMON 1\n"); + // QSA_TOPK=1: the v1 top-k; QSA_ATTN=1|2|3 (default 3): which attention (2 has its own chunking and merge) + const int av = std::getenv("QSA_ATTN") ? std::atoi(std::getenv("QSA_ATTN")) : 1; + const bool v2 = av == 2; + mtl::Pipeline pprep = cx.pipeline(lib, "qsa_prep", mtl::Constants().b(24, std::getenv("QSA_MROPE") != nullptr)), pscore = cx.pipeline(lib, "qsa_score"), + ptopk = cx.pipeline(lib, std::getenv("QSA_TOPK") ? "qsa_topk" : "qsa_topk2"), + pattn = cx.pipeline(lib, av == 1 ? "qsa_attn" : av == 2 ? "qsa_attn2" : "qsa_attn3"), + pmerge = cx.pipeline(lib, v2 ? "qsa_merge2" : "qsa_merge"); + const int MB = MAXC / 4 + 2; + // per-cell inputs for every cell (context + window) + const int NC = P + MAXT; + mtl::Buffer qfull = cx.alloc(4 * (size_t)NC * QH * 2 * QD), kc = cx.alloc(4 * (size_t)NC * QKV * QD), + vc = cx.alloc(4 * (size_t)NC * QKV * QD), iq = cx.alloc(4 * (size_t)NC * IH * ID), + ik = cx.alloc(4 * (size_t)NC * ID), qn = cx.alloc(4 * QD), kn = cx.alloc(4 * QD), iqn = cx.alloc(4 * ID), + ikn = cx.alloc(4 * ID), cs = cx.alloc(8 * (size_t)(MAXC + 16) * 32), qcur = cx.alloc(4 * MAXT * QH * QD), + qidx = cx.alloc(4 * MAXT * IH * ID), kq = cx.alloc((size_t)MAXC * QKV * QD), vq = cx.alloc((size_t)MAXC * QKV * QD), + ksc = cx.alloc(2 * (size_t)MAXC * QKV * 4), vsc = cx.alloc(2 * (size_t)MAXC * QKV * 4), tail = cx.alloc(4 * 4 * ID), + dead = cx.alloc(4 * ID), pooled = cx.alloc(4 * (size_t)MB * ID), scores = cx.alloc(4 * (size_t)MAXT * MB), + ids = cx.alloc(4 * MAXT * BUDGET), part = cx.alloc(4 * (size_t)MAXT * 2 * NCH * QH * (QD + 2)), + xp = cx.alloc(4 * (4 * MAXT * QH * QD + MAXT * QH * QD / 16)); + ref::Rng r(23); + for (size_t i = 0; i < qfull.bytes / 4; ++i) qfull.as()[i] = r.sym() * 2.0f; + for (size_t i = 0; i < kc.bytes / 4; ++i) kc.as()[i] = r.sym() * 2.0f; + for (size_t i = 0; i < vc.bytes / 4; ++i) vc.as()[i] = r.sym(); + for (size_t i = 0; i < iq.bytes / 4; ++i) iq.as()[i] = r.sym(); + for (size_t i = 0; i < ik.bytes / 4; ++i) ik.as()[i] = r.sym(); + for (int i = 0; i < QD; ++i) { qn.as()[i] = 1.0f + 0.2f * r.sym(); kn.as()[i] = 1.0f + 0.2f * r.sym(); } + for (int i = 0; i < ID; ++i) { iqn.as()[i] = 1.0f + 0.2f * r.sym(); ikn.as()[i] = 1.0f + 0.2f * r.sym(); } + const uint32_t pos_base = 0; + // QSA_MROPE=1: qsa_prep with the M-RoPE table (function constant 24): random (t, h, w) positions per cell, the CPU + // reference rotating pair i by row tab[cell * 3 + i % 3]; QSA_MROPE=id: the identity table, checked bitwise + // against the plain kernel + const char* mre = std::getenv("QSA_MROPE"); + const bool mrope = mre != nullptr, mid = mrope && std::string(mre) == "id"; + mtl::Buffer mtab = cx.alloc(4 * 3 * (size_t)(MAXC + 16)); + { + ref::Rng rm(77); + for (int c = 0; c < MAXC + 16; ++c) + for (int k = 0; k < 3; ++k) mtab.as()[c * 3 + k] = mid ? c : (int)(rm.next() % (uint64_t)(MAXC + 16)); + } + for (int p = 0; p < MAXC + 16; ++p) + for (int i = 0; i < 32; ++i) { + const double inv = std::pow(1e7, -2.0 * i / 64.0), th = (double)p * inv; + cs.as()[(p * 32 + i) * 2] = (float)std::cos(th); + cs.as()[(p * 32 + i) * 2 + 1] = (float)std::sin(th); + } + auto prep = [&](uint32_t pos0, int nt) { + QsaArgs a{pos0, pos_base, (uint32_t)MAXC, (uint32_t)nt, 1e-6f}; + mtl::Dispatch d; + d.p = &pprep; + d.args = {mtl::Arg::B(qfull, 4 * (size_t)pos0 * QH * 2 * QD), mtl::Arg::B(kc, 4 * (size_t)pos0 * QKV * QD), + mtl::Arg::B(vc, 4 * (size_t)pos0 * QKV * QD), mtl::Arg::B(iq, 4 * (size_t)pos0 * IH * ID), + mtl::Arg::B(ik, 4 * (size_t)pos0 * ID), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(iqn), + mtl::Arg::B(ikn), mtl::Arg::B(cs), mtl::Arg::B(qcur), mtl::Arg::B(qidx), mtl::Arg::B(kq), + mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), mtl::Arg::B(tail), mtl::Arg::B(dead), + mtl::Arg::B(pooled), mtl::Arg::V(a), mtl::Arg::B(mtab)}; + d.grid[0] = 33; + d.grid[1] = nt; + d.tg[0] = 256; + return d; + }; + auto window = [&](uint32_t pos0, int nt, std::vector& ds) { + QsaArgs a{pos0, pos_base, (uint32_t)MAXC, (uint32_t)nt, 1e-6f}; + ds.push_back(prep(pos0, nt)); + mtl::Dispatch s; + s.p = &pscore; + s.args = {mtl::Arg::B(qidx), mtl::Arg::B(pooled), mtl::Arg::B(dead), mtl::Arg::B(scores), mtl::Arg::V(a)}; + s.grid[0] = ((pos0 + nt) / 4 + 1 + 7) / 8; + s.grid[1] = nt; + s.tg[0] = 256; + ds.push_back(s); + mtl::Dispatch k; + k.p = &ptopk; + k.args = {mtl::Arg::B(scores), mtl::Arg::B(ids), mtl::Arg::V(a)}; + k.grid[0] = nt; + k.tg[0] = 1024; + ds.push_back(k); + mtl::Dispatch at; + at.p = &pattn; + at.args = {mtl::Arg::B(qcur), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), + mtl::Arg::B(ids), mtl::Arg::B(part), mtl::Arg::V(a)}; + at.grid[0] = v2 ? (BUDGET + 31) / 32 : NCH; + at.grid[1] = v2 ? QKV * 2 : QKV; + at.grid[2] = nt; + at.tg[0] = 256; + ds.push_back(at); + mtl::Dispatch m; + m.p = &pmerge; + m.args = {mtl::Arg::B(part), mtl::Arg::B(qfull, 4 * (size_t)pos0 * QH * 2 * QD), mtl::Arg::B(xp), mtl::Arg::V(a)}; + m.grid[0] = QH; + m.grid[1] = nt; + m.tg[0] = 256; + ds.push_back(m); + }; + // the context, one cell at a time + { + std::vector ds; + for (int p = 0; p < P; ++p) ds.push_back(prep(p, 1)); + cx.run(ds); + } + if (bench) { + std::printf("device %s, context %d cells\n", cx.device_name().c_str(), P); + for (int nt : {1, 3}) { + std::vector ds; + window(P, nt, ds); + const char* names[] = {"prep", "score", "topk", "attn", "merge"}; + double total = 0; + for (int i = 0; i < 5; ++i) { + std::vector one(32, ds[i]); + cx.warm(one, 0.05); + std::vector t; + for (int k = 0; k < 7; ++k) t.push_back(cx.run(one) / 32); + std::printf(" ntok %d %-6s %7.2f us\n", nt, names[i], median(t) * 1e6); + total += median(t); + } + std::vector t; + for (int k = 0; k < 7; ++k) t.push_back(cx.run(ds)); + std::printf(" ntok %d window (5 dependent dispatches): %.2f us (sum of parts %.2f)\n", nt, median(t) * 1e6, + total * 1e6); + } + return 0; + } + const int NT = 3; + std::vector ds; + window(P, NT, ds); + // stage by stage, re-reading GPU state in between + cx.run({ds[0]}); + int fails = 0; + // ---- prep: qcur, qidx, the window cells' KV codes and the pooled keys vs CPU + { + Err eq, ei, ekv, ep; + auto normrope = [&](const float* x, const float* g, int w, int absp, bool rope, std::vector& out) { + double ss = 0; + for (int i = 0; i < w; ++i) ss += (double)x[i] * x[i]; + const double rs = 1.0 / std::sqrt(ss / w + 1e-6); + out.assign(w, 0.0); + for (int i = 0; i < w; ++i) out[i] = rs * x[i] * g[i]; + if (rope) + for (int i = 0; i < 32; ++i) { + const int pp = mrope ? mtab.as()[absp * 3 + i % 3] : absp; + const double c = cs.as()[(pp * 32 + i) * 2], s = cs.as()[(pp * 32 + i) * 2 + 1]; + const double a = out[i], b = out[i + 32]; + out[i] = a * c - b * s; + out[i + 32] = a * s + b * c; + } + }; + std::vector o; + for (int t = 0; t < NT; ++t) { + const int cell = P + t; + for (int h = 0; h < QH; ++h) { + normrope(qfull.as() + ((size_t)cell * QH + h) * 2 * QD, qn.as(), QD, cell, true, o); + for (int i = 0; i < QD; ++i) eq.add(qcur.as()[(t * QH + h) * QD + i], o[i]); + } + for (int h = 0; h < IH; ++h) { + normrope(iq.as() + ((size_t)cell * IH + h) * ID, iqn.as(), ID, cell, true, o); + for (int i = 0; i < ID; ++i) ei.add(qidx.as()[(t * IH + h) * ID + i], o[i]); + } + for (int h = 0; h < QKV; ++h) { + normrope(kc.as() + ((size_t)cell * QKV + h) * QD, kn.as(), QD, cell, true, o); + for (int i = 0; i < QD; ++i) { + const size_t row = (size_t)cell * QKV + h; + const double deq = (double)(int8_t)kq.as()[row * QD + i] * h2f(ksc.as()[row * 4 + i / 64]); + ekv.add(deq, o[i]); + } + } + } + // pooled keys of the last few completed blocks + const int nb = (P + NT) / 4; + for (int b = std::max(0, nb - 8); b < nb; ++b) { + std::vector m(ID); + for (int i = 0; i < ID; ++i) + m[i] = (((ik.as()[(4 * b) * ID + i] + ik.as()[(4 * b + 1) * ID + i]) + + ik.as()[(4 * b + 2) * ID + i]) + ik.as()[(4 * b + 3) * ID + i]) * 0.25f; + normrope(m.data(), ikn.as(), ID, 4 * b, true, o); + for (int i = 0; i < ID; ++i) ep.add(pooled.as()[b * ID + i], o[i]); + } + // int8 quantization error is ~1/254 of max|x| per group: the KV check is a quantization tolerance + const bool ok = eq.l2() < 1e-5 && ei.l2() < 1e-5 && ep.l2() < 1e-5 && ekv.l2() < 6e-3; + fails += !ok; + std::printf("prep: q rel L2 %.2e, indexer q %.2e, pooled keys %.2e, K int8 (vs float) %.2e %s\n", eq.l2(), + ei.l2(), ep.l2(), ekv.l2(), ok ? "PASS" : "FAIL"); + } + // ---- score + cx.run({ds[1]}); + { + Err es; + for (int t = 0; t < NT; ++t) { + const int n_kv = P + t + 1, n_bid = n_kv / 4; + for (int b = 0; b <= n_bid; ++b) { + const float* key = b == n_bid ? dead.as() : pooled.as() + b * ID; + double s = 0; + for (int h = 0; h < IH; ++h) { + double d = 0; + for (int i = 0; i < ID; ++i) d += (double)qidx.as()[(t * IH + h) * ID + i] * key[i]; + s += std::max(d, 0.0); + } + if (b == n_bid && (n_kv & 3)) { + // float(s + 1e9) keeps no bits of s: only check that the bias is there + if (!(scores.as()[t * MB + b] >= 1e9f)) es.add(1.0, 0.0); + } else { + es.add(scores.as()[t * MB + b], s); + } + } + } + const bool ok = es.l2() < 1e-5; + fails += !ok; + std::printf("score: rel L2 %.2e (tail +1e9 checked separately) %s\n", es.l2(), ok ? "PASS" : "FAIL"); + } + // ---- top-k (exact, from the GPU's scores) + cx.run({ds[2]}); + { + int bad = 0; + for (int t = 0; t < NT; ++t) { + const int n_kv = P + t + 1, n_bid = n_kv / 4, width = std::min(n_kv, BUDGET); + std::vector cells(n_kv); + for (int c = 0; c < n_kv; ++c) cells[c] = c; + if (n_kv > BUDGET) { + const float* sc = scores.as() + t * MB; + std::stable_sort(cells.begin(), cells.end(), [&](int x, int y) { return sc[x / 4] > sc[y / 4]; }); + cells.resize(width); + std::sort(cells.begin(), cells.end()); + } + (void)n_bid; + for (int i = 0; i < width; ++i) bad += ids.as()[t * BUDGET + i] != cells[i]; + } + fails += bad != 0; + std::printf("top-k: %d of %d ids differ (context %d: selection %s) %s\n", bad, NT * std::min(P + 1, BUDGET), P, + P + 1 > BUDGET ? "active" : "identity", bad ? "FAIL" : "PASS"); + } + // ---- attention + merge (from the GPU's qcur, ids and KV codes) + cx.run({ds[3], ds[4]}); + { + Err eo; + for (int t = 0; t < NT; ++t) { + const int n_kv = P + t + 1, width = std::min(n_kv, BUDGET); + for (int h = 0; h < QH; ++h) { + const int g = h / (QH / QKV); + std::vector s(width); + double mx = -1e300; + for (int c = 0; c < width; ++c) { + const size_t row = (size_t)ids.as()[t * BUDGET + c] * QKV + g; + double d = 0; + for (int i = 0; i < QD; ++i) + d += (double)qcur.as()[(t * QH + h) * QD + i] * (int8_t)kq.as()[row * QD + i] * + h2f(ksc.as()[row * 4 + i / 64]); + s[c] = d / 16.0; + mx = std::max(mx, s[c]); + } + double l = 0; + for (int c = 0; c < width; ++c) l += (s[c] = std::exp(s[c] - mx)); + for (int i = 0; i < QD; ++i) { + double o = 0; + for (int c = 0; c < width; ++c) { + const size_t row = (size_t)ids.as()[t * BUDGET + c] * QKV + g; + o += s[c] * (int8_t)vq.as()[row * QD + i] * h2f(vsc.as()[row * 4 + i / 64]); + } + o /= l; + const double gate = qfull.as()[((size_t)(P + t) * QH + h) * 2 * QD + QD + i]; + eo.add(xp.as()[t * QH * QD + h * QD + i], o / (1.0 + std::exp(-gate))); + } + } + } + const bool ok = eo.l2() < 1e-5; + fails += !ok; + std::printf("attention + gate: rel L2 %.2e %s\n", eo.l2(), ok ? "PASS" : "FAIL"); + } + std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/qsa_prompt_parity.cpp b/edge0-strata/metal/tests/qsa_prompt_parity.cpp new file mode 100644 index 00000000..18e1b3af --- /dev/null +++ b/edge0-strata/metal/tests/qsa_prompt_parity.cpp @@ -0,0 +1,160 @@ +// metal/tests/qsa_prompt_parity.cpp - qsa_attn_mm (qsa_prompt.metal, the matrix units, fp16 operands) against qsa_attn +// (fp32, itself checked by qsa_parity), both followed by qsa_merge: the attention output of T query tokens on random +// int8 caches, within the budget (all cells, causal) and past it (2051 selected cells per token); then the speed. +// qsa_prompt_parity [--bench] +#include "coop_probe.h" +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int QH = 24, QKV = 2, QD = 256, QB = 2051, NCH = 33; +struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; }; +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } +} // namespace + +int main(int argc, char** argv) { + const bool bench = argc > 1 && std::string(argv[1]) == "--bench"; + std::setvbuf(stdout, nullptr, _IOLBF, 0); + const bool noflash = std::getenv("QSA_NOFLASH") != nullptr; + mtl::Context cx; + const int lib = cx.compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\n"); + std::printf("device %s\n", cx.device_name().c_str()); + const int T = bench ? 2048 : 96, MAXC = 40000; + mtl::Buffer kq = cx.alloc((size_t)MAXC * QKV * QD), vq = cx.alloc((size_t)MAXC * QKV * QD), ksc = cx.alloc(2 * (size_t)MAXC * QKV * 4), + vsc = cx.alloc(2 * (size_t)MAXC * QKV * 4), qcur = cx.alloc(4 * (size_t)T * QH * QD), qh = cx.alloc(2 * (size_t)T * QH * QD), + qfull = cx.alloc(4 * (size_t)T * QH * 2 * QD), ids = cx.alloc(4 * (size_t)T * QB), + part = cx.alloc(4 * (size_t)T * NCH * QH * (QD + 2)), xa = cx.alloc(4 * ((size_t)4 * T * QH * QD + (size_t)T * QH * QD / 16)), + xb = cx.alloc(4 * ((size_t)4 * T * QH * QD + (size_t)T * QH * QD / 16)); + ref::Rng r(5); + for (size_t i = 0; i < kq.bytes; ++i) { kq.as()[i] = (int8_t)((int)(r.next() % 255) - 127); vq.as()[i] = (int8_t)((int)(r.next() % 255) - 127); } + for (size_t i = 0; i < ksc.bytes / 2; ++i) { ksc.as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); vsc.as()[i] = ref::f32_to_f16(0.004f + 0.004f * r.uni()); } + for (size_t i = 0; i < (size_t)T * QH * QD; ++i) qcur.as()[i] = r.sym() * 3.0f; + for (size_t i = 0; i < (size_t)T * QH * 2 * QD; ++i) qfull.as()[i] = r.sym(); + mtl::Pipeline pa = cx.pipeline(lib, "qsa_attn"), pm = cx.pipeline(lib, "qsa_merge"), pmm = cx.pipeline(lib, "qsa_attn_mm", mtl::Constants().b(7, coop_k4(cx)).i(8, std::getenv("QSA_STOP") ? std::atoi(std::getenv("QSA_STOP")) : 0)), + pq = cx.pipeline(lib, "qsa_q_f16"), pf = cx.pipeline(lib, "qsa_flash", mtl::Constants().b(7, coop_k4(cx))); + mtl::Buffer xf = cx.alloc(2 * (size_t)T * QH * QD); + std::printf("pipelines: qsa_attn_mm max threads %d, tg mem %d; qsa_flash max threads %d, tg mem %d\n", pmm.max_threads, + pmm.static_tg_mem, pf.max_threads, pf.static_tg_mem); + int fails = 0; + for (uint32_t pos0 : {0u, 30000u}) { + // ids: within the budget the identity; past it 2051 distinct ascending cells < n_kv + for (int t = 0; t < T; ++t) { + const uint32_t n_kv = pos0 + t + 1; + int* o = ids.as() + (size_t)t * QB; + if (n_kv <= (uint32_t)QB) { for (uint32_t c = 0; c < n_kv; ++c) o[c] = (int)c; continue; } + std::set s; + while ((int)s.size() < QB) s.insert((int)(r.next() % n_kv)); + int j = 0; + for (int c : s) o[j++] = c; + } + QsaArgs a{pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f}; + auto D = [&](const mtl::Pipeline& p, std::vector args, uint32_t gx, uint32_t gy, uint32_t gz, uint32_t tg) { + mtl::Dispatch d; d.p = &p; d.args = std::move(args); d.grid[0] = gx; d.grid[1] = gy; d.grid[2] = gz; d.tg[0] = tg; return d; + }; + const uint32_t nq = T * QH * QD; + std::vector ref_ds = { + D(pa, {mtl::Arg::B(qcur), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), mtl::Arg::B(ids), mtl::Arg::B(part), mtl::Arg::V(a)}, NCH, QKV, T, 256), + D(pm, {mtl::Arg::B(part), mtl::Arg::B(qfull), mtl::Arg::B(xa), mtl::Arg::V(a)}, QH, T, 1, 256)}; + std::vector mm_ds = { + D(pq, {mtl::Arg::B(qcur), mtl::Arg::B(qh), mtl::Arg::V(nq)}, (nq + 255) / 256, 1, 1, 256), + D(pmm, {mtl::Arg::B(qh), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), mtl::Arg::B(ids), mtl::Arg::B(part), mtl::Arg::V(a)}, NCH, QKV, T, 32), + D(pm, {mtl::Arg::B(part), mtl::Arg::B(qfull), mtl::Arg::B(xb), mtl::Arg::V(a)}, QH, T, 1, 256)}; + const mtl::Dispatch fl = D(pf, {mtl::Arg::B(qh), mtl::Arg::B(kq), mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), + mtl::Arg::B(ids), mtl::Arg::B(qfull), mtl::Arg::B(xf), mtl::Arg::V(a)}, QKV, T, 1, 128); + if (bench) { + auto timeit = [&](const std::vector& ds) { + cx.warm(ds, 0.3); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run(ds)); + return median(t); + }; + std::printf("pos0 %5u, T %d: qsa_attn %.2f ms, qsa_attn_mm %.2f ms, merge %.2f ms; qsa_flash (no merge) %.2f ms\n", pos0, + T, timeit({ref_ds[0]}) * 1e3, timeit({mm_ds[1]}) * 1e3, timeit({ref_ds[1]}) * 1e3, timeit({fl}) * 1e3); + continue; + } + cx.run(ref_ds); + std::fill(part.as(), part.as() + part.bytes / 4, std::nanf("")); // a kernel that writes nothing shows + cx.run(mm_ds); + if (std::getenv("QSA_DEBUG")) { + // how much of the partials the matrix-unit kernel wrote, and token 5's first chunk (m, l, acc[0..3]) + size_t written = 0, nonfinite = 0; + const float* pp = part.as(); + for (size_t i = 0; i < (size_t)T * NCH * QH * (QD + 2); ++i) { + if (!std::isnan(pp[i])) ++written; + if (!std::isfinite(pp[i]) && !std::isnan(pp[i])) ++nonfinite; + } + const float* q = pp + ((size_t)5 * NCH * QH) * (QD + 2); + std::printf("debug pos0 %u: %zu of %zu partial floats written (%zu inf); t5 h0: m %g l %g acc %g %g %g %g\n", pos0, + written, (size_t)T * NCH * QH * (QD + 2), nonfinite, q[QD], q[QD + 1], q[0], q[1], q[2], q[3]); + } + double se = 0, sr = 0; + for (size_t i = 0; i < (size_t)T * QH * QD; ++i) { + const double e = xb.as()[i] - xa.as()[i]; + se += e * e; + sr += (double)xa.as()[i] * xa.as()[i]; + } + const double rel = std::sqrt(se / sr); + const bool ok = rel < 3e-3; + fails += !ok; + std::printf("pos0 %5u, T %d: qsa_attn_mm vs qsa_attn (after merge) rel L2 %.2e %s\n", pos0, T, rel, ok ? "PASS" : "FAIL"); + if (noflash) continue; + cx.run({fl}); + double sf = 0; + for (size_t i = 0; i < (size_t)T * QH * QD; ++i) { + const double e = ref::f16_to_f32(xf.as()[i]) - xa.as()[i]; + sf += e * e; + } + const double relf = std::sqrt(sf / sr); + const bool okf = relf < 3e-3; + fails += !okf; + std::printf("pos0 %5u, T %d: qsa_flash (fp16 out) vs qsa_attn + merge rel L2 %.2e %s\n", pos0, T, relf, okf ? "PASS" : "FAIL"); + } + // qsa_score_mm vs qsa_score at pos0 30000: the scores of every block <= n_bid + { + const uint32_t pos0 = 30000, MB = MAXC / 4 + 2, nb = (pos0 + T) / 4 + 1; + mtl::Buffer qidx = cx.alloc(4 * (size_t)T * 4 * 128), pooled = cx.alloc(4 * (size_t)MB * 128), dead = cx.alloc(4 * 128), + sa = cx.alloc(4 * (size_t)T * MB), sb = cx.alloc(4 * (size_t)T * MB); + for (size_t i = 0; i < (size_t)T * 4 * 128; ++i) qidx.as()[i] = r.sym(); + for (size_t i = 0; i < (size_t)MB * 128; ++i) pooled.as()[i] = r.sym(); + for (int i = 0; i < 128; ++i) dead.as()[i] = r.sym(); + QsaArgs a{pos0, 0, (uint32_t)MAXC, (uint32_t)T, 1e-6f}; + mtl::Pipeline ps = cx.pipeline(lib, "qsa_score"), psm = cx.pipeline(lib, "qsa_score_mm"); + mtl::Dispatch d1; d1.p = &ps; d1.args = {mtl::Arg::B(qidx), mtl::Arg::B(pooled), mtl::Arg::B(dead), mtl::Arg::B(sa), mtl::Arg::V(a)}; + d1.grid[0] = (nb + 7) / 8; d1.grid[1] = T; d1.tg[0] = 256; + mtl::Dispatch d2; d2.p = &psm; d2.args = {mtl::Arg::B(qidx), mtl::Arg::B(pooled), mtl::Arg::B(dead), mtl::Arg::B(sb), mtl::Arg::V(a)}; + d2.grid[0] = (nb + 63) / 64; d2.grid[1] = (T + 15) / 16; d2.tg[0] = 128; + if (bench) { + auto timeit = [&](const mtl::Dispatch& d) { + cx.warm({d}, 0.3); + std::vector t; + for (int i = 0; i < 5; ++i) t.push_back(cx.run({d})); + return median(t); + }; + std::printf("score pos0 %u, T %d: qsa_score %.2f ms, qsa_score_mm %.2f ms\n", pos0, T, timeit(d1) * 1e3, timeit(d2) * 1e3); + } else { + cx.run({d1}); + cx.run({d2}); + double worst = 0; + for (int t = 0; t < T; ++t) { + const uint32_t nbid = (pos0 + t + 1) / 4; + for (uint32_t b = 0; b <= nbid; ++b) { + const double x = sa.as()[(size_t)t * MB + b], y = sb.as()[(size_t)t * MB + b]; + worst = std::max(worst, std::fabs(x - y) / (std::fabs(x) + 1.0)); // scores are O(10) + } + } + const bool ok = worst < 1e-5; + fails += !ok; + std::printf("score pos0 %u, T %d: qsa_score_mm vs qsa_score worst rel %.2e %s\n", pos0, T, worst, ok ? "PASS" : "FAIL"); + } + } + if (!bench) std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/ref_dump.cpp b/edge0-strata/metal/tests/ref_dump.cpp new file mode 100644 index 00000000..3cce42e6 --- /dev/null +++ b/edge0-strata/metal/tests/ref_dump.cpp @@ -0,0 +1,30 @@ +// metal/tests/ref_dump.cpp - writes random blocks of every type ref_quant knows and their CPU-reference dequant, so +// that the reference itself can be checked against gguf-py (metal/tools/check_ref.py): +// ref_dump -> /.bin (raw blocks, 16 rows x 2560) and /.f32 +#include "ref_quant.h" + +#include +#include +#include + +int main(int argc, char** argv) { + const std::string dir = argc > 1 ? argv[1] : "."; + const int rows = 16, cols = 2560; + for (int t : {ref::Q2_0, ref::Q4_0, ref::Q5_0, ref::Q8_0, ref::Q3_K, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::IQ2_XXS, + ref::IQ2_XS, ref::IQ2_S, ref::IQ3_XXS, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS, ref::BF16, ref::F16}) { + const size_t rb = ref::row_bytes(t, cols); + std::vector raw(rb * rows); + ref::random_blocks(t, raw.data(), rows, cols, 77 + t); + std::vector f((size_t)rows * cols); + ref::dequant(t, raw.data(), f.data(), f.size()); + const std::string n = dir + "/" + ref::type_name(t); + FILE* a = std::fopen((n + ".bin").c_str(), "wb"); + FILE* b = std::fopen((n + ".f32").c_str(), "wb"); + std::fwrite(raw.data(), 1, raw.size(), a); + std::fwrite(f.data(), 4, f.size(), b); + std::fclose(a); + std::fclose(b); + } + std::printf("wrote %s\n", dir.c_str()); + return 0; +} diff --git a/edge0-strata/metal/tests/route_parity.cpp b/edge0-strata/metal/tests/route_parity.cpp new file mode 100644 index 00000000..2f72329a --- /dev/null +++ b/edge0-strata/metal/tests/route_parity.cpp @@ -0,0 +1,90 @@ +// metal/tests/route_parity.cpp - moe_route (softmax + top-10 + the expert plan) against a CPU replica of +// native_router.cu's rule, for 1..8 tokens; logits with deliberate ties. +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int P_NG = 0, P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4, + P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE; +struct MoeArgs { + uint64_t gate_off, up_off, down_off, gu_stride, d_stride; + uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad; +}; +} // namespace + +int main() { + mtl::Context cx; + const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal"}); + mtl::Pipeline p = cx.pipeline(lib, "moe_route"); + mtl::Buffer lg = cx.alloc(4 * 8 * 512), ids = cx.alloc(4 * 80), w = cx.alloc(4 * 80), plan = cx.alloc(4 * P_INTS); + ref::Rng r(9); + int fails = 0; + double worst_w = 0; + for (int trial = 0; trial < 200; ++trial) { + const int nt = 1 + trial % 8; + for (int i = 0; i < nt * 512; ++i) { + float v = r.sym() * 4.0f; + if (trial % 3 == 0) v = std::round(v * 4.0f) / 4.0f; // many exact ties + lg.as()[i] = v; + } + MoeArgs a{}; + a.n_tok = nt; + a.topk = 10; + a.n_expert = 512; + a.gx_gu = 1; + a.gx_dn = 1; + mtl::Dispatch d; + d.p = &p; + d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a)}; + d.tg[0] = 512; + cx.run({d}); + for (int t = 0; t < nt; ++t) { + const float* l = lg.as() + t * 512; + float mx = -INFINITY; + for (int e = 0; e < 512; ++e) mx = std::max(mx, l[e]); + std::vector pr(512); + double s = 0; + for (int e = 0; e < 512; ++e) s += (pr[e] = std::exp((double)l[e] - mx)); + for (int e = 0; e < 512; ++e) pr[e] /= s; + // top 10: largest probability, ties -> lower index (probabilities of equal logits are equal) + std::vector order(512); + for (int e = 0; e < 512; ++e) order[e] = e; + std::stable_sort(order.begin(), order.end(), [&](int x, int y) { return l[x] > l[y]; }); + double ssum = 0; + for (int k = 0; k < 10; ++k) ssum += pr[order[k]]; + for (int k = 0; k < 10; ++k) { + if (ids.as()[t * 10 + k] != order[k]) { ++fails; break; } + const double ref = pr[order[k]] / std::max(ssum, 6.103515625e-5); + worst_w = std::max(worst_w, std::fabs(w.as()[t * 10 + k] - ref) / ref); + } + } + // the plan + const int* P = plan.as(); + std::set distinct(ids.as(), ids.as() + nt * 10); + bool ok = P[P_NG] == (int)distinct.size(); + int gi = 0; + for (int e : distinct) { + if (P[P_GEXP + gi] != e) ok = false; + int j = P[P_GSTART + gi]; + for (int i = 0; i < nt * 10; ++i) + if (ids.as()[i] == e) { + if (P[P_ETOK + j] != i / 10 || P[P_TOKENT + i] != j) ok = false; + ++j; + } + if (j != P[P_GSTART + gi + 1]) ok = false; + ++gi; + } + fails += !ok; + } + std::printf("route parity, 200 windows of 1..8 tokens (a third with exact ties): ids/plan mismatches %d, worst weight " + "rel err %.2e %s\n", fails, worst_w, fails == 0 && worst_w < 2e-6 ? "PASS" : "FAIL"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/rt_bench.cpp b/edge0-strata/metal/tests/rt_bench.cpp new file mode 100644 index 00000000..9fc49ff7 --- /dev/null +++ b/edge0-strata/metal/tests/rt_bench.cpp @@ -0,0 +1,238 @@ +// metal/tests/rt_bench.cpp - the M-series GPU's achievable read bandwidth and the cost of a dispatch. +// +// rt_bench bw [MiB] streaming read kernels over a buffer (default 2048 MiB, far beyond the SLC), per +// threadgroup size and grid size; read+write (copy); small sizes (cache-resident) +// rt_bench dispatch empty / dependent / small kernels: N dispatches in one serial or concurrent encoder, +// with or without barriers, through an indirect command buffer, and one command buffer +// per dispatch (commit + wait round trip) +#include "mtl_rt.h" + +#include +#include +#include +#include +#include +#include + +static double median(std::vector v) { + std::sort(v.begin(), v.end()); + return v[v.size() / 2]; +} + +template static void timed(int warm, int reps, F f, double& med, double& best) { + for (int i = 0; i < warm; ++i) f(); + std::vector t; + for (int i = 0; i < reps; ++i) t.push_back(f()); + med = median(t); + best = *std::min_element(t.begin(), t.end()); +} + +static void bw(mtl::Context& cx, int lib, size_t mib) { + const size_t bytes = mib << 20; + mtl::Buffer src = cx.alloc(bytes); + for (size_t i = 0; i < bytes / 4; ++i) src.as()[i] = (uint32_t)(i * 2654435761u); + mtl::Buffer out = cx.alloc(64u << 20); + const uint32_t n = (uint32_t)(bytes / 16); + mtl::Pipeline pr = cx.pipeline(lib, "k_read"), pu = cx.pipeline(lib, "k_read_u4"), + pc = cx.pipeline(lib, "k_read_chunk"); + std::printf("buffer %zu MiB; k_read max threads %d, simd width %d\n", mib, pr.max_threads, pr.simd_width); + std::printf("%-12s %6s %8s %10s %10s\n", "kernel", "tg", "groups", "GB/s med", "GB/s best"); + double best_all = 0; + std::string best_cfg; + for (const mtl::Pipeline* p : {&pr, &pu}) { + for (uint32_t tg : {64u, 128u, 256u, 512u, 1024u}) { + if ((int)tg > p->max_threads) continue; + for (uint32_t groups : {80u, 160u, 320u, 640u, 1280u, 2560u, 5120u}) { + if ((size_t)groups * tg * 4 > out.bytes) continue; + mtl::Dispatch d; + d.p = p; + d.args = {mtl::Arg::B(src), mtl::Arg::B(out), mtl::Arg::V(n)}; + d.grid[0] = groups; + d.tg[0] = tg; + double med, best; + timed(2, 7, [&] { return cx.run({d}); }, med, best); + const double gm = bytes / med / 1e9, gb = bytes / best / 1e9; + std::printf("%-12s %6u %8u %10.1f %10.1f\n", p == &pr ? "read" : "read_u4", tg, groups, gm, gb); + if (gm > best_all) { + best_all = gm; + best_cfg = std::string(p == &pr ? "read" : "read_u4") + " tg " + std::to_string(tg) + " groups " + + std::to_string(groups); + } + } + } + } + // contiguous chunk per threadgroup + for (uint32_t tg : {256u, 512u, 1024u}) { + for (uint32_t chunk_kib : {16u, 64u, 256u, 1024u}) { + const uint32_t chunk = chunk_kib * 1024 / 16; + const uint32_t groups = (n + chunk - 1) / chunk; + if ((size_t)groups * tg * 4 > out.bytes) continue; + mtl::Dispatch d; + d.p = &pc; + d.args = {mtl::Arg::B(src), mtl::Arg::B(out), mtl::Arg::V(n), mtl::Arg::V(chunk)}; + d.grid[0] = groups; + d.tg[0] = tg; + double med, best; + timed(2, 7, [&] { return cx.run({d}); }, med, best); + const double gm = bytes / med / 1e9; + std::printf("chunk %4uK %6u %8u %10.1f %10.1f\n", chunk_kib, tg, groups, gm, bytes / best / 1e9); + if (gm > best_all) { + best_all = gm; + best_cfg = "chunk " + std::to_string(chunk_kib) + "K tg " + std::to_string(tg); + } + } + } + std::printf("BEST read: %.1f GB/s (%s)\n", best_all, best_cfg.c_str()); + + // read + write + { + mtl::Buffer dst = cx.alloc(bytes / 2); + mtl::Pipeline pk = cx.pipeline(lib, "k_copy"); + const uint32_t nn = (uint32_t)(bytes / 2 / 16); + double bestc = 0; + for (uint32_t tg : {256u, 1024u}) + for (uint32_t groups : {640u, 2560u}) { + mtl::Dispatch d; + d.p = &pk; + d.args = {mtl::Arg::B(src), mtl::Arg::B(dst), mtl::Arg::V(nn)}; + d.grid[0] = groups; + d.tg[0] = tg; + double med, best; + timed(2, 7, [&] { return cx.run({d}); }, med, best); + bestc = std::max(bestc, (double)bytes / med / 1e9); + } + std::printf("BEST copy (read+write, both counted): %.1f GB/s\n", bestc); + cx.free(dst); + } + // cache-resident sizes: the same kernel on a small prefix + for (size_t kib : {1024ul, 4096ul, 16384ul, 32768ul, 65536ul, 131072ul, 262144ul}) { + if (kib * 1024 > bytes) break; + const uint32_t nn = (uint32_t)(kib * 1024 / 16); + double bestg = 0; + for (uint32_t groups : {160u, 640u, 2560u}) { + mtl::Dispatch d; + d.p = &pu; + d.args = {mtl::Arg::B(src), mtl::Arg::B(out), mtl::Arg::V(nn)}; + d.grid[0] = groups; + d.tg[0] = 256; + // many repetitions in one command buffer, so the launch is amortised + const int reps = (int)std::max(1, (256ul << 10) / kib * 4); + double med, best; + timed(1, 5, [&] { return cx.run({d}, reps) / reps; }, med, best); + bestg = std::max(bestg, kib * 1024.0 / med / 1e9); + } + std::printf("resident %7zu KiB: %.1f GB/s\n", kib, bestg); + } + cx.free(src); + cx.free(out); +} + +// A read of `kib` from DRAM per dispatch (each dispatch a different, uncached chunk of a 4 GiB buffer), dependent +// dispatches in one serial encoder: the bound for a kernel that streams that much once. +static void bwsize(mtl::Context& cx, int lib) { + const size_t big = 4ull << 30; + mtl::Buffer src = cx.alloc(big), out = cx.alloc(64u << 20); + mtl::Pipeline pr = cx.pipeline(lib, "k_read"); + std::printf("%10s %8s %10s %10s\n", "chunk KiB", "tg x grp", "us/disp", "GB/s"); + for (size_t kib : {1152ul, 2304ul, 4500ul, 9000ul, 13500ul, 18000ul, 36000ul, 72000ul}) { + const size_t bytes = kib * 1024; + const int R = (int)std::min(400, big / bytes); + double bestt = 1e9; + uint32_t bg = 0, btg = 0; + for (uint32_t tg : {256u, 1024u}) + for (uint32_t groups : {80u, 160u, 320u, 640u, 1280u}) { + std::vector ds; + const uint32_t n = (uint32_t)(bytes / 16); + for (int r = 0; r < R; ++r) { + mtl::Dispatch d; + d.p = ≺ + d.args = {mtl::Arg::B(src, (size_t)r * ((big / R) & ~4095ull)), mtl::Arg::B(out), mtl::Arg::V(n)}; + d.grid[0] = groups; + d.tg[0] = tg; + ds.push_back(d); + } + double med, best; + timed(1, 5, [&] { return cx.run(ds) / R; }, med, best); + if (med < bestt) { bestt = med; bg = groups; btg = tg; } + } + std::printf("%10zu %4ux%-4u %10.2f %10.1f\n", kib, btg, bg, bestt * 1e6, bytes / bestt / 1e9); + } + cx.free(src); + cx.free(out); +} + +static void dispatch(mtl::Context& cx, int lib) { + mtl::Buffer buf = cx.alloc(1 << 20), x = cx.alloc(4 * 2560), y = cx.alloc(4 * 2560); + mtl::Pipeline pe = cx.pipeline(lib, "k_empty"), pch = cx.pipeline(lib, "k_chain"), ps = cx.pipeline(lib, "k_small"); + auto mk = [&](const mtl::Pipeline& p, int kind) { + mtl::Dispatch d; + d.p = &p; + if (kind == 2) { + d.args = {mtl::Arg::B(x), mtl::Arg::B(y)}; + d.grid[0] = 10; + d.tg[0] = 256; + } else { + d.args = {mtl::Arg::B(buf)}; + d.grid[0] = 1; + d.tg[0] = 32; + } + return d; + }; + const int N = 2000; + struct K { const char* name; const mtl::Pipeline* p; int kind; } ks[] = { + {"empty", &pe, 0}, {"chain", &pch, 1}, {"small", &ps, 2}}; + std::printf("%-8s %-36s %10s\n", "kernel", "mode", "us/dispatch"); + for (const K& k : ks) { + std::vector ds(N, mk(*k.p, k.kind)); + double med, best; + timed(2, 7, [&] { return cx.run(ds, 1, false, false); }, med, best); + std::printf("%-8s %-36s %10.2f\n", k.name, "serial encoder", med / N * 1e6); + timed(2, 7, [&] { return cx.run(ds, 1, true, false); }, med, best); + std::printf("%-8s %-36s %10.2f\n", k.name, "concurrent encoder, no barrier", med / N * 1e6); + timed(2, 7, [&] { return cx.run(ds, 1, true, true); }, med, best); + std::printf("%-8s %-36s %10.2f\n", k.name, "concurrent encoder + barrier each", med / N * 1e6); + timed(2, 7, [&] { return cx.run(ds, 1, false, true); }, med, best); + std::printf("%-8s %-36s %10.2f\n", k.name, "serial encoder + barrier each", med / N * 1e6); + timed(2, 7, [&] { return cx.run_icb(ds, 1, false); }, med, best); + std::printf("%-8s %-36s %10.2f\n", k.name, "ICB, no barrier", med / N * 1e6); + timed(2, 7, [&] { return cx.run_icb(ds, 1, true); }, med, best); + std::printf("%-8s %-36s %10.2f\n", k.name, "ICB + barrier each", med / N * 1e6); + } + // command buffer round trips + { + std::vector one(1, mk(pe, 0)); + double gpu = 0; + const int R = 300; + cx.run_wall(one, 20); + const double wall = cx.run_wall(one, R, &gpu); + std::printf("%-8s %-36s %10.2f (GPU %.2f us)\n", "empty", "1 cmdbuf/dispatch, commit+wait wall", wall / R * 1e6, + gpu / R * 1e6); + std::vector d100(100, mk(pe, 0)); + const double wall2 = cx.run_wall(d100, 100, &gpu); + std::printf("%-8s %-36s %10.2f (GPU %.2f us)\n", "empty", "100 dispatches/cmdbuf, wall per cb", wall2 / 100 * 1e6, + gpu / 100 * 1e6); + } + // the fixed cost of one command buffer as seen by the GPU timer (1 dispatch) + { + std::vector one(1, mk(pe, 0)); + double med, best; + timed(5, 21, [&] { return cx.run(one); }, med, best); + std::printf("%-8s %-36s %10.2f\n", "empty", "single dispatch cmdbuf GPU time", med * 1e6); + } + cx.free(buf); + cx.free(x); + cx.free(y); +} + +int main(int argc, char** argv) { + mtl::Context cx; + std::printf("device: %s apple9 %d apple10 %d metal4 %d working set %.1f GiB\n", cx.device_name().c_str(), + cx.has_family_apple(9), cx.has_family_apple(10), cx.supports_metal4(), + cx.recommended_working_set() / 1073741824.0); + const int lib = cx.compile({"rt_bench.metal"}, "#define NO_GGML_COMMON 1\n"); + const std::string mode = argc > 1 ? argv[1] : "all"; + if (mode == "bw" || mode == "all") bw(cx, lib, argc > 2 ? std::strtoul(argv[2], nullptr, 10) : 2048); + if (mode == "dispatch" || mode == "all") dispatch(cx, lib); + if (mode == "bwsize" || mode == "all") bwsize(cx, lib); + return 0; +} diff --git a/edge0-strata/metal/tests/runtime_hello.cpp b/edge0-strata/metal/tests/runtime_hello.cpp new file mode 100644 index 00000000..eee9e5f1 --- /dev/null +++ b/edge0-strata/metal/tests/runtime_hello.cpp @@ -0,0 +1,162 @@ +// metal/tests/runtime_hello.cpp - the Metal runtime end to end, with the shared src/core/graph.cpp built unchanged: +// 1. gpu::malloc_device / host_alloc (mapped) / memcpy_async / a kernel through metal::dispatch / events / D2H +// 2. the SHARED host file src/core/graph.cpp (CapturedGraph / GraphRegistry), compiled unchanged for Metal: +// capture a 3-kernel "window", replay it 5 times, wait with wait_ms (event query polling), check the data path +// is live (each replay reads the buffers' current contents, as a CUDA graph does) +// 3. host_device_pointer: a kernel writes a mapped host buffer that the host reads without a copy +// 4. the cost of replaying a 1,000-dispatch window +#include "strata/core/graph.hpp" +#include "strata/gpu/metal.hpp" +#include "strata/gpu/runtime.hpp" + +#include +#include +#include +#include +#include + +namespace gpu = strata::gpu; + +static const char* kSrc = R"MSL( +#include +using namespace metal; +struct Axpy { device const float* x; device float* y; float a; uint n; }; +kernel void axpy(constant Axpy& p [[buffer(0)]], uint i [[thread_position_in_grid]]) { + if (i < p.n) p.y[i] = p.a * p.x[i] + p.y[i]; +} +struct Scale { device float* y; device const float* s; uint n; }; // the factor is DATA (a device buffer), not an argument +kernel void scale(constant Scale& p [[buffer(0)]], uint i [[thread_position_in_grid]]) { + if (i < p.n) p.y[i] *= p.s[0]; +} +)MSL"; +struct Axpy { uint64_t x, y; float a; uint32_t n; }; +struct Scale { uint64_t y, s; uint32_t n; }; + +static int g_fail = 0; +#define CHECK(c, ...) do { if (!(c)) { std::printf("FAIL %s:%d: ", __FILE__, __LINE__); std::printf(__VA_ARGS__); std::printf("\n"); ++g_fail; } } while (0) + +static double now_ms() { + return std::chrono::duration(std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +int main() { + std::string err; + MTL::ComputePipelineState* p_axpy = gpu::metal::pipeline("hello", kSrc, "axpy", err); + MTL::ComputePipelineState* p_scale = gpu::metal::pipeline("hello", kSrc, "scale", err); + CHECK(p_axpy && p_scale, "pipelines: %s", err.c_str()); + if (!p_axpy || !p_scale) return 1; + const uint32_t N = 1 << 16; + const gpu::metal::Grid groups{(N + 255) / 256, 1, 1}, threads{256, 1, 1}; + + gpu::Stream s = nullptr, s2 = nullptr; + CHECK(gpu::stream_create(&s, gpu::kStreamNonBlocking) == gpu::kSuccess, "stream"); + CHECK(gpu::stream_create(&s2, gpu::kStreamNonBlocking) == gpu::kSuccess, "stream 2"); + float *dx = nullptr, *dy = nullptr, *ds = nullptr, *hx = nullptr; + CHECK(gpu::malloc_device((void**) &dx, N * 4) == gpu::kSuccess, "malloc x"); + CHECK(gpu::malloc_device((void**) &dy, N * 4) == gpu::kSuccess, "malloc y"); + CHECK(gpu::malloc_device((void**) &ds, 4) == gpu::kSuccess, "malloc s"); + CHECK(gpu::host_alloc((void**) &hx, N * 4, gpu::kHostAllocMapped) == gpu::kSuccess, "host_alloc"); + CHECK(gpu::host_view(dx) != nullptr && gpu::host_view(dx) != (void*) dx, "a device pointer is a GPU address, not the CPU view"); + + // ---- 1. copies, a kernel, an event across two streams, D2H + std::vector x(N), y(N, 1.0f), got(N); + for (uint32_t i = 0; i < N; ++i) x[i] = (float) (i % 1000); + std::memcpy(hx, x.data(), N * 4); + gpu::memcpy_async(dx, hx, N * 4, gpu::kHostToDevice, s); // mapped host buffer -> device: a blit + gpu::memcpy_async(dy, y.data(), N * 4, gpu::kHostToDevice, s); // pageable -> device + Axpy a{(uint64_t) dx, (uint64_t) dy, 2.0f, N}; + gpu::metal::dispatch(s, p_axpy, &a, sizeof a, groups, threads); + gpu::Event ev = nullptr; + gpu::event_create(&ev, gpu::kEventDisableTiming); + gpu::event_record(ev, s); + gpu::stream_wait_event(s2, ev); // s2's kernel runs after s's + Axpy a2{(uint64_t) dx, (uint64_t) dy, 1.0f, N}; + gpu::metal::dispatch(s2, p_axpy, &a2, sizeof a2, groups, threads); + gpu::memcpy_async(got.data(), dy, N * 4, gpu::kDeviceToHost, s2); + gpu::stream_synchronize(s2); + int bad = 0; + for (uint32_t i = 0; i < N; ++i) bad += got[i] != 3.0f * x[i] + 1.0f; + CHECK(bad == 0, "step 1: %d wrong values", bad); + std::printf("1. copies + kernel + cross-stream event + D2H: %s\n", bad ? "FAIL" : "OK"); + + // ---- 2. the shared CapturedGraph from src/core/graph.cpp + { + float one = 1.0f; + gpu::memcpy_sync(ds, &one, 4, gpu::kHostToDevice); + strata::core::CapturedGraph g; + CHECK(g.begin(s, err), "begin: %s", err.c_str()); + Axpy w{(uint64_t) dx, (uint64_t) dy, 1.0f, N}; + Scale sc{(uint64_t) dy, (uint64_t) ds, N}; + gpu::metal::dispatch(s, p_axpy, &w, sizeof w, groups, threads); + gpu::metal::dispatch(s, p_scale, &sc, sizeof sc, groups, threads); + gpu::metal::dispatch(s, p_axpy, &w, sizeof w, groups, threads); + CHECK(g.end(s, err), "end: %s", err.c_str()); + CHECK(g.nodes() == 3, "nodes %zu", g.nodes()); + // host model of y + std::vector m(N); + for (uint32_t i = 0; i < N; ++i) m[i] = got[i]; + bool all = true; + for (int rep = 0; rep < 5; ++rep) { + const float f = rep % 2 ? 0.5f : 2.0f; // the per-step value changes as DATA between replays + float* sv = (float*) gpu::host_view(ds); + gpu::stream_synchronize(s); + *sv = f; + CHECK(g.launch(s, err), "launch: %s", err.c_str()); + CHECK(g.wait_ms(2000), "wait_ms timed out"); + for (uint32_t i = 0; i < N; ++i) m[i] = ((m[i] + x[i]) * f) + x[i]; + } + gpu::memcpy_sync(got.data(), dy, N * 4, gpu::kDeviceToHost); + int wrong = 0; + for (uint32_t i = 0; i < N; ++i) wrong += got[i] != m[i]; + all = wrong == 0; + CHECK(all, "graph replays: %d wrong values", wrong); + std::printf("2. shared CapturedGraph (src/core/graph.cpp): 3-kernel window, 5 replays with live data: %s\n", all ? "OK" : "FAIL"); + + strata::core::GraphRegistry reg(s); + int bodies = 0; + for (int k = 0; k < 4; ++k) + CHECK(reg.record(strata::core::LayerType::GDN, 1, [&] { ++bodies; gpu::metal::dispatch(s, p_axpy, &w, sizeof w, groups, threads); }, err), + "record: %s", err.c_str()); + CHECK(bodies == 1 && reg.size() == 1, "registry captured %d bodies", bodies); + CHECK(reg.launch(strata::core::LayerType::GDN, 1, 2000, err), "registry launch: %s", err.c_str()); + std::printf(" GraphRegistry: body captured once over 4 records, launch+wait: %s\n", bodies == 1 ? "OK" : "FAIL"); + } + + // ---- 3. a mapped host buffer written by a kernel, read by the host with no copy + { + float* dview = nullptr; + CHECK(gpu::host_device_pointer((void**) &dview, hx) == gpu::kSuccess, "host_device_pointer"); + std::memset(hx, 0, N * 4); + Axpy w{(uint64_t) dx, (uint64_t) dview, 1.0f, N}; + gpu::metal::dispatch(s, p_axpy, &w, sizeof w, groups, threads); + gpu::stream_synchronize(s); + int wrong = 0; + for (uint32_t i = 0; i < N; ++i) wrong += hx[i] != x[i]; + CHECK(wrong == 0, "mapped: %d wrong", wrong); + std::printf("3. kernel -> mapped host buffer, host reads it in place: %s\n", wrong ? "FAIL" : "OK"); + } + + // ---- 4. replay cost of a 1,000-dispatch window + { + strata::core::CapturedGraph g; + Axpy w{(uint64_t) dx, (uint64_t) dy, 0.0f, 256}; + g.begin(s, err); + for (int k = 0; k < 1000; ++k) gpu::metal::dispatch(s, p_axpy, &w, sizeof w, gpu::metal::Grid{1, 1, 1}, threads); + g.end(s, err); + for (int rep = 0; rep < 5; ++rep) { + const double t0 = now_ms(); + g.launch(s, err); + const double t1 = now_ms(); + g.wait_ms(5000); + std::printf("4. 1000-dispatch window: replay (encode+commit) %.3f ms, until complete %.3f ms\n", t1 - t0, now_ms() - t0); + } + } + const auto st = gpu::metal::stats(); + std::printf("stats: %llu dispatches, %llu command buffers, %llu blits, %llu host syncs\n", (unsigned long long) st.dispatches, + (unsigned long long) st.command_buffers, (unsigned long long) st.blits, (unsigned long long) st.host_syncs); + gpu::free_device(dx); gpu::free_device(dy); gpu::free_device(ds); gpu::free_host(hx); + gpu::event_destroy(ev); + gpu::stream_destroy(s); gpu::stream_destroy(s2); + std::printf(g_fail ? "RESULT: %d FAILED\n" : "RESULT: all passed\n", g_fail); + return g_fail ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/sample_parity.cpp b/edge0-strata/metal/tests/sample_parity.cpp new file mode 100644 index 00000000..049439af --- /dev/null +++ b/edge0-strata/metal/tests/sample_parity.cpp @@ -0,0 +1,61 @@ +// metal/tests/sample_parity.cpp - greedy argmax (sample.metal) vs the CPU over 248,320 logits with planted ties, +// and its time. +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include + +int main() { + constexpr uint32_t V = 248320, PARTS = 120; + mtl::Context cx; + const int lib = cx.compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n"); + mtl::Pipeline pp = cx.pipeline(lib, "argmax_part"), pf = cx.pipeline(lib, "argmax_final"); + const int MAXT = 8; + mtl::Buffer lg = cx.alloc(4 * (size_t)V * MAXT), pv = cx.alloc(4 * PARTS * MAXT), pi = cx.alloc(4 * PARTS * MAXT), + out = cx.alloc(4 * MAXT); + ref::Rng r(5); + int fails = 0; + struct A { uint32_t n_vocab, parts; } a{V, PARTS}; + for (int trial = 0; trial < 20; ++trial) { + const int nt = 1 + trial % MAXT; + for (size_t i = 0; i < (size_t)V * nt; ++i) lg.as()[i] = std::round(r.sym() * 1000.0f) / 100.0f; + for (int t = 0; t < nt; ++t) { + // plant a tie at the maximum + const size_t i1 = r.next() % V, i2 = r.next() % V; + lg.as()[t * V + i1] = 20.0f; + lg.as()[t * V + i2] = 20.0f; + } + mtl::Dispatch d1, d2; + d1.p = &pp; + d1.args = {mtl::Arg::B(lg), mtl::Arg::B(pv), mtl::Arg::B(pi), mtl::Arg::V(a)}; + d1.grid[0] = PARTS; + d1.grid[1] = nt; + d1.tg[0] = 256; + d2.p = &pf; + d2.args = {mtl::Arg::B(pv), mtl::Arg::B(pi), mtl::Arg::B(out), mtl::Arg::V(a)}; + d2.grid[0] = nt; + d2.tg[0] = 256; + cx.run({d1, d2}); + for (int t = 0; t < nt; ++t) { + const float* l = lg.as() + (size_t)t * V; + int best = 0; + for (uint32_t i = 1; i < V; ++i) + if (l[i] > l[best]) best = (int)i; + fails += out.as()[t] != best; + } + if (trial == 19) { + std::vector ds; + for (int k = 0; k < 64; ++k) { ds.push_back(d1); ds.push_back(d2); } + cx.warm(ds, 0.05); + std::vector ts; + for (int k = 0; k < 7; ++k) ts.push_back(cx.run(ds) / 64); + std::sort(ts.begin(), ts.end()); + std::printf("argmax %d token(s) over %u logits: %.2f us (2 dispatches)\n", nt, V, ts[3] * 1e6); + } + } + std::printf("argmax parity, 20 windows with planted ties: %d wrong %s\n", fails, fails ? "FAIL" : "PASS"); + return fails ? 1 : 0; +} diff --git a/edge0-strata/metal/tests/serve_probe.py b/edge0-strata/metal/tests/serve_probe.py new file mode 100644 index 00000000..5d50ff5b --- /dev/null +++ b/edge0-strata/metal/tests/serve_probe.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""Metal serve regression probe (owner round 9, G6): one request of the first N token ids of a prompt through serve/server.py on the +config in STRATA_CFG, with its default flags (expert cache auto, prefill auto). 95, 569, 1500 and 8213 ids used to end in +"prefill: relayout needs borrowed buffers and a chunk of at most 0" or in a copy into device allocations that were given back. + + STRATA_REPO= STRATA_CFG= STRATA_PROMPTS= python3 serve_probe.py EXE 95:agent 8 + +Needs the Gyro-S (or any Flash-Next) model on the Mac; it is not part of the unit tests.""" +import os +import sys +import threading +import time +from pathlib import Path + +REPO = Path(os.environ["STRATA_REPO"]) +sys.path.insert(0, str(REPO / "tools")) +sys.path.insert(0, str(REPO)) +import strata_tokenizer as ST # noqa: E402 +from serve.server import StrataEngine, child_env # noqa: E402 + +EXE = sys.argv[1] +NAME = sys.argv[2] if len(sys.argv) > 2 else "agent" +MAX_NEW = int(sys.argv[3]) if len(sys.argv) > 3 else 8 +S = os.environ.get("STRATA_OUT", os.path.expanduser("~/strata-mac/g023/out")) +cfg = json.load(open(os.environ["STRATA_CFG"], encoding="utf-8-sig")) +tp = Path(cfg["tokenizer"]) +vocab = json.loads((tp / "vocab.json").read_text(encoding="utf-8")) +toks = [None] * len(vocab) +for t, i in vocab.items(): + toks[i] = t +tok = ST.Tokenizer(toks, (tp / "merges.txt").read_text(encoding="utf-8").split("\n"), + json.loads((tp / "token_type.json").read_text())) +if NAME.split(":")[0].isdigit(): # "4096" / "4096:agent": the first N ids of the agent prompt + parts = NAME.split(":") + n = int(parts[0]) + src = parts[1] if len(parts) > 1 else "agent" + raw = (Path(os.environ.get("STRATA_PROMPTS", os.path.expanduser("~/owner-t9/bundle/prompts"))) / f"{src}.prompt.ids").read_text() + try: + ids = json.loads(raw)[:n] + except Exception: + ids = [int(x) for x in raw.replace(",", " ").split()][:n] +else: + raw = (Path(os.environ.get("STRATA_PROMPTS", os.path.expanduser("~/owner-t9/bundle/prompts"))) / f"{NAME}.prompt.ids").read_text() + try: + ids = json.loads(raw) + except Exception: + ids = [int(x) for x in raw.replace(",", " ").split()] +args = list(cfg["args"]) + os.environ.get("STRATA_EXTRA_ARGS", "").split() +log = os.path.join(S, f"probe-{NAME}.log") +t0 = time.time() +eng = StrataEngine(EXE, args, cwd=cfg.get("cwd"), log=log, env=child_env(cfg)) +print(f"READY after {time.time() - t0:.1f} s, prompt {len(ids)} tokens", flush=True) +t1 = time.time() +out = [t for t in eng.generate(ids, MAX_NEW, {"temperature": 0}, threading.Event()) if t is not None] +last = dict(eng.last) +print(f"{NAME}: {len(out)} out in {time.time() - t1:.1f} s; prompt {last.get('prompt_tokens')} tok in " + f"{last.get('prompt_ms')} ms = {last.get('prompt_tokens', 0) / (last.get('prompt_ms', 1) / 1000):.1f} tok/s; " + f"decode {last.get('decode_ms')} ms = {len(out) / (last['decode_ms'] / 1000) if last.get('decode_ms') else 0:.2f} tok/s; " + f"drafts {last.get('drafts_accepted')}/{last.get('drafts_offered')}", flush=True) +print("text:", tok.decode(out)[:120].replace("\n", " "), flush=True) +eng.proc.stdin.close() +try: + eng.proc.wait(60) +except Exception: + eng.proc.kill() +print("PROBE_DONE", flush=True) diff --git a/edge0-strata/metal/tests/token_chain.cpp b/edge0-strata/metal/tests/token_chain.cpp new file mode 100644 index 00000000..a47c7630 --- /dev/null +++ b/edge0-strata/metal/tests/token_chain.cpp @@ -0,0 +1,661 @@ +// metal/tests/token_chain.cpp - one decode step of the whole model, every kernel of metal/ on the REAL weights of a +// GGUF (mmapped, wrapped no-copy as one MTLBuffer), each layer with its own tensor formats, all dispatches in one +// command buffer: the GPU time of a decode token (or a verify window) with these kernels, dispatch boundaries +// included. Not a numerical check (the activations are synthetic): a timing of the real byte stream. +// +// token_chain [--ntok N] [--ctx P] [--reps R] [--skip-qsa] [--only CLASS] +// +// Left out: the PLE path (one table row and two small projections at one layer) and the MTP draft. +#include "gguf_lite.h" +#include "mtl_rt.h" +#include "ref_quant.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int N = 2560, HC = 4, LR = 320, NFF = 640, TOPK = 10, NEXP = 512, MAXT = 8; +constexpr int P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4, + P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE, P_IGU = 4, P_IDN = 8; +struct GemvArgs { uint32_t n_in, n_out, n_tok, row_bytes, x_stride, y_stride; }; +struct GrArgs { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; }; +struct GrArgs3 { uint32_t n_embd, hc, hc_lr, n_tok, apply; float eps; uint32_t rb_down, rb_up, comb, topk; }; +struct MoeArgs { + uint64_t gate_off, up_off, down_off, gu_stride, d_stride; + uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad; +}; +struct GdnArgs { uint32_t h_k, h_v, channels, commit; float eps; uint32_t plane; }; +struct QsaArgs { uint32_t pos0, pos_base, max_cells, n_tok; float eps; }; +struct GlueArgs { uint32_t n, n_tok, n_embd, topk; }; +struct ArgArgs { uint32_t n_vocab, parts; }; + +int xkind(int t) { + switch (t) { + case ref::Q2_0: return 1; + case ref::Q4_K: case ref::Q5_K: return 2; + case ref::Q6_K: case ref::Q3_K: return 3; + default: return 0; + } +} +double median(std::vector v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; } + +// a prepared activation: 4 planes of n_tok rows of n, then the per-16 sums +struct Act { + mtl::Buffer b; + int n = 0; + size_t plane_bytes() const { return (size_t)4 * MAXT * n; } +}; + +} // namespace + +int main(int argc, char** argv) { + if (argc < 2) { std::fprintf(stderr, "usage: token_chain [--ntok N] [--ctx P] [--reps R]\n"); return 2; } + int ntok = 1, ctx = 1000, reps = 10; + bool skip_qsa = false, graph = false; + // the HC read: Q8_0 copies through the 3-dispatch read by default (the owner's decision, as STRATA_HC_Q8 on gfx12); + // STRATA_HC_Q8=0 selects the exact BF16 path. --hc 3: BF16, 3 dispatches; 2: BF16, 2 dispatches; q8: Q8_0 copies, + // 2 dispatches; q8_3: Q8_0 copies, 3 dispatches + const char* hcq8 = std::getenv("STRATA_HC_Q8"); + std::string hc_mode = hcq8 && std::string(hcq8) == "0" ? "3" : "q8_3"; + bool fuse = false; // the GDN y prep in the step; no QSA score/top-k within the selection budget + bool fold_comb = false; // the MoE combine folded into the next 2-dispatch HC read (measured slower) + std::string only, skip; + for (int i = 2; i < argc; ++i) { + std::string a = argv[i]; + if (a == "--ntok") ntok = std::atoi(argv[++i]); + else if (a == "--ctx") ctx = std::atoi(argv[++i]); + else if (a == "--reps") reps = std::atoi(argv[++i]); + else if (a == "--skip-qsa") skip_qsa = true; + else if (a == "--graph") graph = true; + else if (a == "--hc") hc_mode = argv[++i]; + else if (a == "--fuse") fuse = true; + else if (a == "--fold-comb") fold_comb = true; + else if (a == "--only") only = argv[++i]; + else if (a == "--skip") skip = argv[++i]; + } + gguf::File f; + std::string err; + if (!f.open(argv[1], err)) { std::fprintf(stderr, "%s\n", err.c_str()); return 2; } + mtl::Context cx; + std::printf("device %s; %s: %zu tensors, %.1f GB; %d token(s), context %d\n", cx.device_name().c_str(), + argv[1], f.tensors.size(), f.file_bytes / 1e9, ntok, ctx); + // fault the file in (page cache), then wrap it + { + const double t0 = mtl::now_s(); + madvise(f.map, f.map_bytes, MADV_WILLNEED); + volatile uint64_t s = 0; + for (size_t o = 0; o < f.file_bytes; o += 16384) s += ((const uint8_t*)f.map)[o]; + std::printf("file touched in %.1f s\n", mtl::now_s() - t0); + } + mtl::Buffer W = cx.wrap(f.map, f.map_bytes); + auto T = [&](const std::string& name) -> const gguf::Tensor& { + const gguf::Tensor* t = f.find(name); + if (!t) { std::fprintf(stderr, "missing tensor %s\n", name.c_str()); std::exit(2); } + return *t; + }; + + const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal", "gemv.metal", "gr.metal", "gr2.metal", "glue.metal", + "router.metal"}); + const int lgdn = cx.compile({"gdn.metal"}, "#define NO_GGML_COMMON 1\n"); + const int lqsa = cx.compile({"qsa.metal", "qsa2.metal"}, "#define NO_GGML_COMMON 1\n"); + const int lsmp = cx.compile({"sample.metal"}, "#define NO_GGML_COMMON 1\n"); + std::map pcache; + + // ---- buffers + const int MAXC = std::max(4096, (ctx + 64) / 4 * 4 + 64), MB = MAXC / 4 + 2; + mtl::Buffer R = cx.alloc(4 * MAXT * HC * N), bo = cx.alloc(4 * MAXT * N), inj[2] = {cx.alloc(4 * MAXT * HC), cx.alloc(4 * MAXT * HC)}, + xn = cx.alloc(4 * MAXT * HC * N), lo = cx.alloc(4 * MAXT * LR), qkv = cx.alloc(4 * MAXT * 10240), + z = cx.alloc(4 * MAXT * 6144), al = cx.alloc(4 * MAXT * 48), be = cx.alloc(4 * MAXT * 48), + yg = cx.alloc(4 * MAXT * 6144), logits = cx.alloc(4 * MAXT * 512), ids = cx.alloc(4 * MAXT * TOPK), + wts = cx.alloc(4 * MAXT * TOPK), plan = cx.alloc(4 * P_INTS), h = cx.alloc(4 * MAXT * TOPK * NFF), + ye = cx.alloc(4 * MAXT * TOPK * N), gsh = cx.alloc(4 * MAXT * NFF), ush = cx.alloc(4 * MAXT * NFF), + ysh = cx.alloc(4 * MAXT * N), sg = cx.alloc(4 * MAXT * 4), qfull = cx.alloc(4 * MAXT * 12288), + kcur = cx.alloc(4 * MAXT * 512), vcur = cx.alloc(4 * MAXT * 512), iqr = cx.alloc(4 * MAXT * 512), + ikr = cx.alloc(4 * MAXT * 128), qcur = cx.alloc(4 * MAXT * 24 * 256), qidx = cx.alloc(4 * MAXT * 4 * 128), + scores = cx.alloc(4 * (size_t)MAXT * MB), qids = cx.alloc(4 * MAXT * 2051), + part = cx.alloc(4 * (size_t)MAXT * 65 * 24 * 258), cs = cx.alloc(8 * (size_t)(MAXC + 16) * 32), + lm = cx.alloc(4 * (size_t)MAXT * 248320), apv = cx.alloc(4 * MAXT * 120), api = cx.alloc(4 * MAXT * 120), + tokout = cx.alloc(4 * MAXT); + auto act = [&](int n) { Act a; a.n = n; a.b = cx.alloc(4 * ((size_t)4 * MAXT * n + (size_t)MAXT * n / 16)); return a; }; + Act XA = act(N), XF = act(N), XY = act(6144), XQ = act(6144), XH = act(NFF); + for (size_t i = 0; i < R.bytes / 4; ++i) R.as()[i] = 0.01f * (float)((i * 2654435761u) % 200) - 1.0f; + for (int p = 0; p < MAXC + 16; ++p) + for (int i = 0; i < 32; ++i) { + const double th = (double)p * std::pow(1e7, -2.0 * i / 64.0); + cs.as()[(p * 32 + i) * 2] = (float)std::cos(th); + cs.as()[(p * 32 + i) * 2 + 1] = (float)std::sin(th); + } + // per-layer recurrent state and QSA caches + std::vector gst, ghist, ghist2, kq, vq, ksc, vsc, tail, dead, pooled; + for (int l = 0; l < 48; ++l) { + if (l % 4 != 3) { + gst.push_back(cx.alloc(4 * 128 * 48 * 128)); + ghist.push_back(cx.alloc(4 * 10240 * 3)); + ghist2.push_back(cx.alloc(4 * 10240 * 3)); + } else { + kq.push_back(cx.alloc((size_t)MAXC * 2 * 256)); + vq.push_back(cx.alloc((size_t)MAXC * 2 * 256)); + ksc.push_back(cx.alloc(2 * (size_t)MAXC * 2 * 4)); + vsc.push_back(cx.alloc(2 * (size_t)MAXC * 2 * 4)); + tail.push_back(cx.alloc(4 * 4 * 128)); + dead.push_back(cx.alloc(4 * 128)); + pooled.push_back(cx.alloc(4 * (size_t)MB * 128)); + } + } + for (int t = 0; t < MAXT; ++t) + for (int c = 0; c < 2051; ++c) qids.as()[t * 2051 + c] = c; + for (auto& b : vsc) for (size_t i = 0; i < b.bytes / 2; ++i) b.as()[i] = 0x2000; // small fp16 scales + for (auto& b : ksc) for (size_t i = 0; i < b.bytes / 2; ++i) b.as()[i] = 0x2000; + + std::vector ds; + std::vector cls; // class of each dispatch (for --only) + std::map bytes_by; + bool nb = false; // the next push is independent of the previous dispatch (no barrier before it) + auto push = [&](const mtl::Dispatch& d0, const std::string& c, double bytes = 0) { + mtl::Dispatch d = d0; + d.barrier = !nb; + nb = false; + if (!only.empty() && c != only) return; + if (!skip.empty() && c == skip) return; + ds.push_back(d); + cls.push_back(c); + bytes_by[c] += bytes; + }; + // pipelines whose names do not match the function (cached by key): create by function name first + auto Pn = [&](int l, const std::string& fn, const std::string& key, const mtl::Constants& c) -> const mtl::Pipeline* { + auto it = pcache.find(key); + if (it == pcache.end()) it = pcache.emplace(key, cx.pipeline(l, fn, c)).first; + return &it->second; + }; + // ---- SLC prefetch (--prefetch): the next HC read's weights touched concurrently with the first projection GEMV + // after the current read, so the latency-bound read finds them in the SLC + const bool prefetch = std::getenv("CHAIN_PREFETCH") != nullptr; + const bool albe = std::getenv("CHAIN_ALBE") != nullptr; // ssm_alpha + ssm_beta as one gemv_cat + // the router GEMV + moe_route as one router_route (default; CHAIN_ROUTE=0: separate) + const bool routef = !(std::getenv("CHAIN_ROUTE") && std::string(std::getenv("CHAIN_ROUTE")) == "0"); + const bool combf = std::getenv("CHAIN_COMB") != nullptr; // the MoE combine folded into the next gr_norm + bool comb_pending = false; + // routed gate/up: the i-quant grids from constant memory (QUANT_CTAB, function constant 23; bitwise; default, + // CHAIN_CTAB=0: the threadgroup copy) + const bool ctabf = !(std::getenv("CHAIN_CTAB") && std::string(std::getenv("CHAIN_CTAB")) == "0"); + mtl::Buffer rcounter = cx.alloc(16); + struct Range { const mtl::Buffer* b; size_t off, bytes; }; + std::vector pending; + mtl::Buffer sink = cx.alloc(16); + const mtl::Pipeline* ptouch = nullptr; + auto push_prefetch = [&]() { + for (const Range& r : pending) { + const uint32_t n16 = (uint32_t)(r.bytes / 16); + mtl::Dispatch d; + d.p = ptouch; + d.args = {mtl::Arg::B(*r.b, r.off), mtl::Arg::B(sink), mtl::Arg::V(n16)}; + d.grid[0] = (n16 + 1023) / 1024; + d.tg[0] = 256; + nb = true; + push(d, "prefetch", (double)r.bytes); + } + pending.clear(); + }; + // ---- GEMV by tensor, with a tile from the measured table + auto gemv2 = [&](const gguf::Tensor& t, const Act& x, const mtl::Buffer& y, size_t y_stride, const std::string& c, + int ty_override = -1) { + const int n_in = (int)t.shape[0], n_out = t.shape.size() > 1 ? (int)t.shape[1] : 1; + const int ty = ty_override >= 0 ? ty_override : t.type; + int L = 16, Rr = 1, S = 4, K = 1; + if (n_in >= 6144) { L = 32; Rr = 1; S = 8; K = 2; } + else if (n_in == 640) { L = 8; Rr = 1; S = 4; K = 1; } + else if (n_out <= 640 || ty == ref::Q2_0 || ty == ref::BF16) { L = 32; Rr = 1; S = 4; K = 1; } + const std::string fn = "gemv_" + std::to_string(ty) + "_" + std::to_string(L) + "_" + std::to_string(Rr) + "_" + + std::to_string(S) + "_" + std::to_string(K); + const mtl::Pipeline* p = Pn(lib, fn, fn + "#" + std::to_string(ntok), mtl::Constants().i(0, ntok)); + GemvArgs a{(uint32_t)n_in, (uint32_t)n_out, (uint32_t)ntok, (uint32_t)ref::row_bytes(t.type, n_in), (uint32_t)n_in, + (uint32_t)y_stride}; + mtl::Dispatch d; + d.p = p; + d.args = {mtl::Arg::B(W, t.offset), mtl::Arg::B(x.b, x.plane_bytes() * xkind(ty)), mtl::Arg::B(x.b, x.plane_bytes() * 4), + mtl::Arg::B(y), mtl::Arg::V(a)}; + const int rpt = (S / K) * (32 / L) * Rr; + d.grid[0] = (n_out + rpt - 1) / rpt; + d.tg[0] = S * 32; + push(d, c, (double)t.bytes); + if (!pending.empty()) push_prefetch(); + }; + const mtl::Pipeline* pnorm = Pn(lib, "gr_norm", "gr_norm", {}); + const mtl::Pipeline* pnormc = Pn(lib, "gr_norm_comb", "gr_norm_comb", {}); + ptouch = Pn(lib, "slc_touch", "slc_touch", {}); + const mtl::Pipeline* pdown = Pn(lib, "gr_down_8_8", "gr_down_8_8", mtl::Constants().i(1, ntok)); + const mtl::Pipeline* pup = Pn(lib, "gr_up", "gr_up", mtl::Constants().i(1, ntok)); + const mtl::Pipeline* pdownq = Pn(lib, "gr_down_q8_8_8", "gr_down_q8_8_8", mtl::Constants().i(1, ntok)); + const mtl::Pipeline* pupq = Pn(lib, "gr_up_q8", "gr_up_q8", mtl::Constants().i(1, ntok)); + const mtl::Pipeline* pprep = Pn(lib, "xprep", "xprep", {}); + const mtl::Pipeline* proute = Pn(lib, "moe_route", "moe_route", {}); + const mtl::Pipeline* pswi = Pn(lib, "swiglu_prep", "swiglu_prep", {}); + const mtl::Pipeline* pcomb = Pn(lib, "moe_combine_sh", "moe_combine_sh", {}); + const mtl::Pipeline* pgdn = Pn(lgdn, "gdn_step", "gdn_step", mtl::Constants().i(2, ntok)); + const mtl::Pipeline* qprep = Pn(lqsa, "qsa_prep", "qsa_prep", {}); + const mtl::Pipeline* qscore = Pn(lqsa, "qsa_score", "qsa_score", {}); + const int qav = std::getenv("QSA_ATTN") ? std::atoi(std::getenv("QSA_ATTN")) : 1; + const bool qv2 = qav == 2; + const mtl::Pipeline* qtopk = Pn(lqsa, std::getenv("QSA_TOPK") ? "qsa_topk" : "qsa_topk2", "qsa_topk", {}); + const mtl::Pipeline* qattn = Pn(lqsa, qav == 1 ? "qsa_attn" : qv2 ? "qsa_attn2" : "qsa_attn3", "qsa_attn", {}); + const mtl::Pipeline* qmerge = Pn(lqsa, qv2 ? "qsa_merge2" : "qsa_merge", "qsa_merge", {}); + const mtl::Pipeline* pam1 = Pn(lsmp, "argmax_part", "argmax_part", {}); + const mtl::Pipeline* pam2 = Pn(lsmp, "argmax_final", "argmax_final", {}); + + // the 2-dispatch read: residual and partial-dot ping-pong buffers, rs; the Q8_0 copies of the HC matrices + mtl::Buffer Rb[2] = {R, cx.alloc(4 * MAXT * HC * N)}, Db[2] = {cx.alloc(4 * MAXT * (LR + HC) * HC), cx.alloc(4 * MAXT * (LR + HC) * HC)}, + rsb = cx.alloc(4 * MAXT * HC); + int cur_r = 0, cur_d = 0; + std::map q8off; + mtl::Buffer Q8; + if (hc_mode == "q8" || hc_mode == "q8_3") { + size_t tot = 0; + for (const auto& kv : f.tensors) { + const gguf::Tensor& t = kv.second; + if (t.type != ref::BF16 || t.name.find("hc_") == std::string::npos || t.name.find("_norm") != std::string::npos) continue; + uint64_t ne = 1; + for (uint64_t d : t.shape) ne *= d; + q8off[t.name] = tot; + tot += ne / 32 * 34; + } + Q8 = cx.alloc((tot + 4095) / 4096 * 4096); + const double t0 = mtl::now_s(); + std::vector tmp; + for (const auto& kv : q8off) { + const gguf::Tensor& t = *f.find(kv.first); + uint64_t ne = 1; + for (uint64_t d : t.shape) ne *= d; + tmp.resize(ne); + ref::dequant(ref::BF16, (const uint8_t*)f.map + t.offset, tmp.data(), ne); + ref::quantize_q8_0(tmp.data(), ne, Q8.as() + kv.second); + } + std::printf("HC Q8_0 copies: %zu tensors, %.1f MB, quantized in %.1f s\n", q8off.size(), tot / 1e6, mtl::now_s() - t0); + } + int cur_inj = 0; + auto gr_read = [&](const std::string& pre, bool apply, const Act& out, bool final_) { + if (hc_mode != "3" && hc_mode != "q8_3") { + const bool q8 = hc_mode == "q8"; + const gguf::Tensor& tn = T(pre + "_norm.weight"); + const gguf::Tensor& tdn = T(pre + "_down.weight"); + const gguf::Tensor& tup = T(pre + "_up.weight"); + const gguf::Tensor& tin = final_ ? tdn : T(pre + "_inject.weight"); + const int tt = std::min(ntok, 3); + const std::string dn = std::string("gr_down3_") + (q8 ? "8_" : "30_") + std::to_string(tt) + "_2"; + const std::string up = std::string("gr_up3_") + (q8 ? "8" : "30"); + // after an MoE (the next layer's mixer read, the final read) the combine is folded in here + const bool comb = fold_comb && apply && pre.find("hc_ffn") == std::string::npos; + GrArgs3 a{N, HC, LR, MAXT, apply ? 1u : 0u, 1e-6f, (uint32_t)(q8 ? HC * N / 32 * 34 : HC * N * 2), + (uint32_t)(q8 ? LR / 32 * 34 : LR * 2), comb ? 1u : 0u, (uint32_t)TOPK}; + const mtl::Buffer& WB = q8 ? Q8 : W; + const size_t od = q8 ? q8off.at(tdn.name) : tdn.offset, oi = q8 ? q8off.at(tin.name) : tin.offset, + ou = q8 ? q8off.at(tup.name) : tup.offset; + mtl::Dispatch d; + d.p = Pn(lib, dn, dn + "#" + std::to_string(ntok), mtl::Constants().i(1, ntok)); + d.args = {mtl::Arg::B(WB, od), mtl::Arg::B(WB, oi), mtl::Arg::B(Rb[cur_r]), mtl::Arg::B(Rb[1 - cur_r]), + mtl::Arg::B(bo), mtl::Arg::B(Db[cur_d]), mtl::Arg::B(W, tn.offset), mtl::Arg::B(Db[1 - cur_d]), + mtl::Arg::B(rsb), mtl::Arg::V(a), mtl::Arg::B(plan), mtl::Arg::B(ye), mtl::Arg::B(ysh), mtl::Arg::B(sg)}; + d.grid[0] = (LR + HC + 15) / 16; + d.grid[1] = HC; + d.tg[0] = 256; + const double wb = q8 ? (tdn.bytes + (final_ ? 0 : tin.bytes)) * 34.0 / 64.0 : (double)tdn.bytes + (final_ ? 0 : tin.bytes); + push(d, "hc", wb); + cur_r = 1 - cur_r; + cur_d = 1 - cur_d; + mtl::Dispatch u; + u.p = Pn(lib, up, up + "#" + std::to_string(ntok), mtl::Constants().i(1, ntok)); + u.args = {mtl::Arg::B(WB, ou), mtl::Arg::B(Db[cur_d]), mtl::Arg::B(Rb[cur_r]), mtl::Arg::B(out.b), mtl::Arg::V(a), + mtl::Arg::B(W, tn.offset), mtl::Arg::B(rsb)}; + u.grid[0] = N / 16; + u.tg[0] = 256; + push(u, "hc", q8 ? tup.bytes * 34.0 / 64.0 : (double)tup.bytes); + return; + } + const gguf::Tensor& tn = T(pre + "_norm.weight"); + const gguf::Tensor& tdn = T(pre + "_down.weight"); + const gguf::Tensor& tup = T(pre + "_up.weight"); + const gguf::Tensor& tin = final_ ? tdn : T(pre + "_inject.weight"); + GrArgs a{N, HC, LR, (uint32_t)ntok, apply ? 1u : 0u, 1e-6f}; + const bool q83 = hc_mode == "q8_3"; + const mtl::Buffer& WB = q83 ? Q8 : W; + const size_t od = q83 ? q8off.at(tdn.name) : tdn.offset, oi = q83 ? q8off.at(tin.name) : tin.offset, + ou = q83 ? q8off.at(tup.name) : tup.offset; + mtl::Dispatch d1; + d1.p = pnorm; + d1.args = {mtl::Arg::B(R), mtl::Arg::B(bo), mtl::Arg::B(inj[cur_inj]), mtl::Arg::B(W, tn.offset), mtl::Arg::B(xn), + mtl::Arg::V(a)}; + if (comb_pending) { + // the MoE combine folded into this norm (gr_norm_comb; bitwise moe_combine_sh + gr_norm) + const uint32_t tk = TOPK; + d1.p = pnormc; + d1.args = {mtl::Arg::B(R), mtl::Arg::B(ye), mtl::Arg::B(inj[cur_inj]), mtl::Arg::B(W, tn.offset), mtl::Arg::B(xn), + mtl::Arg::V(a), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sg), mtl::Arg::V(tk)}; + comb_pending = false; + } + d1.grid[0] = HC; + d1.grid[1] = ntok; + d1.tg[0] = 256; + push(d1, "hc", 0); + mtl::Dispatch d2; + d2.p = q83 ? pdownq : pdown; + d2.args = {mtl::Arg::B(WB, od), mtl::Arg::B(WB, oi), mtl::Arg::B(xn), mtl::Arg::B(lo), + mtl::Arg::B(inj[1 - cur_inj]), mtl::Arg::V(a)}; + d2.grid[0] = LR + HC; + d2.tg[0] = 256; + push(d2, "hc", (double)tdn.bytes + (final_ ? 0 : tin.bytes)); + mtl::Dispatch d3; + d3.p = q83 ? pupq : pup; + d3.args = {mtl::Arg::B(WB, ou), mtl::Arg::B(lo), mtl::Arg::B(xn), mtl::Arg::B(out.b), mtl::Arg::V(a)}; + d3.grid[0] = N / 16; + d3.tg[0] = 256; + push(d3, "hc", (double)tup.bytes); + cur_inj = 1 - cur_inj; + if (prefetch && !final_) { + // the next read: blk.l.hc_attn -> blk.l.hc_ffn -> blk.(l+1).hc_attn ... -> output_hc + std::string nx; + const size_t at = pre.find(".hc_attn"); + if (at != std::string::npos) nx = pre.substr(0, at) + ".hc_ffn"; + else { + const int l = std::atoi(pre.c_str() + 4); + nx = l == 47 ? std::string("output_hc") : "blk." + std::to_string(l + 1) + ".hc_attn"; + } + const bool fin = nx == "output_hc"; + for (const std::string& part : fin ? std::vector{"_down", "_up"} + : std::vector{"_down", "_inject", "_up"}) { + const gguf::Tensor& tt = T(nx + part + ".weight"); + if (q83) { + uint64_t ne = 1; + for (uint64_t dd : tt.shape) ne *= dd; + pending.push_back({&Q8, q8off.at(tt.name), (size_t)(ne / 32 * 34) / 16 * 16}); + } else { + pending.push_back({&W, tt.offset, (size_t)tt.bytes / 16 * 16}); + } + } + } + }; + auto xprep = [&](const mtl::Buffer& x, const Act& out, int n, const std::string& c) { + mtl::Dispatch d; + d.p = pprep; + const uint32_t nn = n, nt = MAXT; + d.args = {mtl::Arg::B(x), mtl::Arg::B(out.b), mtl::Arg::V(nn), mtl::Arg::V(nt)}; + d.tg[0] = 256; + d.grid[0] = (n / 16 * MAXT + 255) / 256; + push(d, c, 0); + }; + // Every prepared activation has planes of MAXT rows (Act::plane_bytes); the producers that take a token count + // for the plane stride (gr_up, qsa_merge, swiglu_prep, xprep) are given MAXT there. + + int gi = 0, qi = 0; + for (int l = 0; l < 48; ++l) { + const std::string b = "blk." + std::to_string(l) + "."; + // ---- mixer + gr_read(b + "hc_attn", l > 0, XA, false); + if (l % 4 != 3) { + gemv2(T(b + "attn_qkv.weight"), XA, qkv, 10240, "gdn proj"); + nb = true; + gemv2(T(b + "attn_gate.weight"), XA, z, 6144, "gdn proj"); + nb = true; + if (albe) { + // ssm_alpha + ssm_beta in one dispatch (gemv_cat; bitwise the two GEMVs) + const gguf::Tensor &ta = T(b + "ssm_alpha.weight"), &tb = T(b + "ssm_beta.weight"); + const std::string fn = "gemvc_" + std::to_string(ta.type) + "_32_1_4"; + mtl::Dispatch d; + d.p = Pn(lib, fn, fn + "#" + std::to_string(ntok), mtl::Constants().i(0, ntok)); + GemvArgs ga{(uint32_t)N, 96, (uint32_t)ntok, (uint32_t)ref::row_bytes(ta.type, N), (uint32_t)N, 48}; + const uint32_t sp = 48; + d.args = {mtl::Arg::B(W, ta.offset), mtl::Arg::B(XA.b, XA.plane_bytes() * xkind(ta.type)), + mtl::Arg::B(XA.b, XA.plane_bytes() * 4), mtl::Arg::B(al), mtl::Arg::V(ga), mtl::Arg::B(W, tb.offset), + mtl::Arg::B(be), mtl::Arg::V(sp)}; + d.grid[0] = 96 / 4; + d.tg[0] = 128; + push(d, "gdn proj", (double)ta.bytes + tb.bytes); + } else { + gemv2(T(b + "ssm_alpha.weight"), XA, al, 48, "gdn proj"); + nb = true; + gemv2(T(b + "ssm_beta.weight"), XA, be, 48, "gdn proj"); + } + GdnArgs a{16, 48, 10240, 1, 1e-6f, fuse ? (uint32_t)(MAXT * 6144) : 0u}; + mtl::Dispatch d; + d.p = pgdn; + d.args = {mtl::Arg::B(gst[gi]), mtl::Arg::B(ghist[gi]), mtl::Arg::B(qkv), mtl::Arg::B(al), mtl::Arg::B(be), + mtl::Arg::B(z), mtl::Arg::B(W, T(b + "ssm_conv1d.weight").offset), mtl::Arg::B(W, T(b + "ssm_dt.bias").offset), + mtl::Arg::B(W, T(b + "ssm_a").offset), mtl::Arg::B(W, T(b + "ssm_norm.weight").offset), mtl::Arg::B(fuse ? XY.b : yg), + mtl::Arg::V(a), mtl::Arg::B(ghist2[gi])}; + d.grid[0] = 48; + d.tg[0] = 512; + push(d, "gdn step", 2.0 * 4 * 128 * 48 * 128); + if (!fuse) xprep(yg, XY, 6144, "gdn proj"); + gemv2(T(b + "ssm_out.weight"), XY, bo, N, "gdn proj"); + ++gi; + } else if (!skip_qsa) { + gemv2(T(b + "attn_q.weight"), XA, qfull, 12288, "qsa proj"); + nb = true; + gemv2(T(b + "attn_k.weight"), XA, kcur, 512, "qsa proj"); + nb = true; + gemv2(T(b + "attn_v.weight"), XA, vcur, 512, "qsa proj"); + nb = true; + gemv2(T(b + "indexer.q_proj.weight"), XA, iqr, 512, "qsa proj"); + nb = true; + gemv2(T(b + "indexer.k_proj.weight"), XA, ikr, 128, "qsa proj"); + QsaArgs a{(uint32_t)ctx, 0, (uint32_t)MAXC, (uint32_t)ntok, 1e-6f}; + mtl::Dispatch d; + d.p = qprep; + d.args = {mtl::Arg::B(qfull), mtl::Arg::B(kcur), mtl::Arg::B(vcur), mtl::Arg::B(iqr), mtl::Arg::B(ikr), + mtl::Arg::B(W, T(b + "attn_q_norm.weight").offset), mtl::Arg::B(W, T(b + "attn_k_norm.weight").offset), + mtl::Arg::B(W, T(b + "indexer.q_norm.weight").offset), mtl::Arg::B(W, T(b + "indexer.k_norm.weight").offset), + mtl::Arg::B(cs), mtl::Arg::B(qcur), mtl::Arg::B(qidx), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]), + mtl::Arg::B(ksc[qi]), mtl::Arg::B(vsc[qi]), mtl::Arg::B(tail[qi]), mtl::Arg::B(dead[qi]), + mtl::Arg::B(pooled[qi]), mtl::Arg::V(a)}; + d.grid[0] = 33; + d.grid[1] = ntok; + d.tg[0] = 256; + push(d, "qsa attn", 0); + const bool sel = !(fuse && ctx + ntok <= 2051); // within the budget the selection is the identity + mtl::Dispatch s; + s.p = qscore; + s.args = {mtl::Arg::B(qidx), mtl::Arg::B(pooled[qi]), mtl::Arg::B(dead[qi]), mtl::Arg::B(scores), mtl::Arg::V(a)}; + s.grid[0] = ((ctx + ntok) / 4 + 1 + 7) / 8; + s.grid[1] = ntok; + s.tg[0] = 256; + if (sel) push(s, "qsa attn", 0); + mtl::Dispatch k; + k.p = qtopk; + k.args = {mtl::Arg::B(scores), mtl::Arg::B(qids), mtl::Arg::V(a)}; + k.grid[0] = ntok; + k.tg[0] = 1024; + if (sel) push(k, "qsa attn", 0); + mtl::Dispatch at; + at.p = qattn; + at.args = {mtl::Arg::B(qcur), mtl::Arg::B(kq[qi]), mtl::Arg::B(vq[qi]), mtl::Arg::B(ksc[qi]), mtl::Arg::B(vsc[qi]), + mtl::Arg::B(qids), mtl::Arg::B(part), mtl::Arg::V(a)}; + at.grid[0] = qv2 ? 65 : 33; + at.grid[1] = qv2 ? 4 : 2; + at.grid[2] = ntok; + at.tg[0] = 256; + push(at, "qsa attn", 2.0 * 2 * 512 * std::min(ctx, 2051)); + mtl::Dispatch m; + m.p = qmerge; + QsaArgs am = a; + am.n_tok = MAXT; // the plane stride (see above) + m.args = {mtl::Arg::B(part), mtl::Arg::B(qfull), mtl::Arg::B(XQ.b), mtl::Arg::V(am)}; + m.grid[0] = 24; + m.grid[1] = ntok; + m.tg[0] = 256; + push(m, "qsa attn", 0); + gemv2(T(b + "attn_output.weight"), XQ, bo, N, "qsa proj"); + ++qi; + } + // ---- MoE + gr_read(b + "hc_ffn", true, XF, false); + if (!routef) { + gemv2(T(b + "ffn_gate_inp.weight"), XF, logits, 512, "router"); + nb = true; + } + const bool pair = fuse && std::getenv("CHAIN_PAIR"); // measured slower (+0.11 ms): opt-in + if (pair) { + // the shared expert's gate + up + SwiGLU in one dispatch, h raw into XH's raw plane + const gguf::Tensor &tgs = T(b + "ffn_gate_shexp.weight"), &tus = T(b + "ffn_up_shexp.weight"); + const std::string fn = "gemv_pair_" + std::to_string(tgs.type) + "_" + std::to_string(tus.type); + mtl::Dispatch d; + d.p = Pn(lib, fn, fn + "#" + std::to_string(ntok), mtl::Constants().i(0, ntok)); + GemvArgs ga{(uint32_t)N, (uint32_t)NFF, (uint32_t)ntok, (uint32_t)ref::row_bytes(tgs.type, N), (uint32_t)N, + (uint32_t)NFF}; + const uint32_t rbu = (uint32_t)ref::row_bytes(tus.type, N); + d.args = {mtl::Arg::B(W, tgs.offset), mtl::Arg::B(XF.b, XF.plane_bytes() * xkind(tgs.type)), + mtl::Arg::B(W, tus.offset), mtl::Arg::B(XF.b, XF.plane_bytes() * xkind(tus.type)), + mtl::Arg::B(XF.b, XF.plane_bytes() * 4), mtl::Arg::B(XH.b), mtl::Arg::V(ga), mtl::Arg::V(rbu)}; + d.grid[0] = NFF / 4; + d.tg[0] = 128; + push(d, "shared", (double)tgs.bytes + tus.bytes); + } else { + gemv2(T(b + "ffn_gate_shexp.weight"), XF, gsh, NFF, "shared"); + nb = true; + gemv2(T(b + "ffn_up_shexp.weight"), XF, ush, NFF, "shared"); + } + nb = true; + gemv2(T(b + "ffn_gate_inp_shexp.weight"), XF, sg, 1, "shared"); + { + MoeArgs a{}; + const gguf::Tensor &tg_ = T(b + "ffn_gate_exps.weight"), &tu = T(b + "ffn_up_exps.weight"), + &td = T(b + "ffn_down_exps.weight"); + a.gate_off = tg_.offset; + a.up_off = tu.offset; + a.down_off = td.offset; + a.gu_row = (uint32_t)ref::row_bytes(tg_.type, N); + a.d_row = (uint32_t)ref::row_bytes(td.type, NFF); + a.gu_stride = (uint64_t)a.gu_row * NFF; + a.d_stride = (uint64_t)a.d_row * N; + a.n_embd = N; + a.n_ff = NFF; + a.n_tok = ntok; + a.topk = TOPK; + a.n_expert = NEXP; + // tiles (lanes per row, rows per lane, simdgroups): CHAIN_GU / CHAIN_DN / CHAIN_DN4 override + const bool dq2 = td.type == ref::Q2_0; + int gl = 32, gr = 1, gs = 4, dl = dq2 ? 8 : 4, dr = dq2 ? 1 : 2, ds_ = 4; + if (const char* e = std::getenv("CHAIN_GU")) std::sscanf(e, "%d,%d,%d", &gl, &gr, &gs); + if (const char* e = std::getenv(dq2 ? "CHAIN_DN" : "CHAIN_DN4")) std::sscanf(e, "%d,%d,%d", &dl, &dr, &ds_); + a.gx_gu = NFF / (gs * (32 / gl) * gr); + a.gx_dn = N / (ds_ * (32 / dl) * dr); + if (routef) { + // the router GEMV + moe_route in one dispatch (router_route; bitwise), concurrent with the shared GEMVs + const gguf::Tensor& tr = T(b + "ffn_gate_inp.weight"); + mtl::Dispatch r; + r.p = Pn(lib, "router_route_30", "router_route#" + std::to_string(ntok), mtl::Constants().i(0, ntok)); + GemvArgs ga{(uint32_t)N, 512, (uint32_t)ntok, (uint32_t)ref::row_bytes(tr.type, N), (uint32_t)N, 512}; + r.args = {mtl::Arg::B(W, tr.offset), mtl::Arg::B(XF.b, 0), mtl::Arg::B(XF.b, XF.plane_bytes() * 4), + mtl::Arg::B(logits), mtl::Arg::V(ga), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan), + mtl::Arg::V(a), mtl::Arg::B(rcounter)}; + r.grid[0] = 512 / 16; + r.tg[0] = 512; + nb = true; + push(r, "router", (double)tr.bytes); + } else { + mtl::Dispatch r; + r.p = proute; + r.args = {mtl::Arg::B(logits), mtl::Arg::B(ids), mtl::Arg::B(wts), mtl::Arg::B(plan), mtl::Arg::V(a)}; + r.tg[0] = 512; + push(r, "router", 0); + } + { + GlueArgs ga{NFF, MAXT, N, TOPK}; + mtl::Dispatch d; + d.p = pswi; + d.args = {mtl::Arg::B(gsh), mtl::Arg::B(ush), mtl::Arg::B(XH.b), mtl::Arg::V(ga)}; + d.grid[0] = NFF / 16; + d.grid[1] = ntok; + d.tg[0] = 16; + nb = true; + if (!pair) push(d, "shared", 0); + nb = false; + } + const std::string gun = "moe_gu_" + std::to_string(tg_.type) + "_" + std::to_string(gl) + "_" + + std::to_string(gr) + "_" + std::to_string(gs); + const std::string dnn = "moe_down_" + std::to_string(dq2 ? 142 : td.type) + "_" + std::to_string(dl) + "_" + + std::to_string(dr) + "_" + std::to_string(ds_); + mtl::Dispatch g; + g.p = Pn(lib, gun, gun + (ctabf ? "#ct" : ""), mtl::Constants().b(23, ctabf)); + g.args = {mtl::Arg::B(W), mtl::Arg::B(XF.b, XF.plane_bytes() * xkind(tg_.type)), mtl::Arg::B(h), mtl::Arg::B(plan), + mtl::Arg::V(a), mtl::Arg::B(XF.b, XF.plane_bytes() * 4)}; + g.tg[0] = 32 * gs; + g.indirect = &plan; + g.ind_off = P_IGU * 4; + push(g, "experts", 10.0 * ntok * 2 * a.gu_stride); + nb = true; + { + const gguf::Tensor& tds = T(b + "ffn_down_shexp.weight"); + // with --fuse h is raw: a Q2_0 down reads it through the raw-activation decoder (type 142) + gemv2(tds, XH, ysh, N, "shared", pair && tds.type == ref::Q2_0 ? 142 : -1); + } + mtl::Dispatch dd; + dd.p = Pn(lib, dnn, dnn, {}); + dd.args = {mtl::Arg::B(W), mtl::Arg::B(h), mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(h)}; + dd.tg[0] = 32 * ds_; + dd.indirect = &plan; + dd.ind_off = P_IDN * 4; + push(dd, "experts", 10.0 * ntok * a.d_stride); + } + { + GlueArgs a{N, (uint32_t)ntok, N, TOPK}; + mtl::Dispatch d; + d.p = pcomb; + d.args = {mtl::Arg::B(ye), mtl::Arg::B(plan), mtl::Arg::B(ysh), mtl::Arg::B(sg), mtl::Arg::B(bo), mtl::Arg::V(a)}; + d.grid[0] = (ntok * N + 255) / 256; + d.tg[0] = 256; + if (combf && (hc_mode == "3" || hc_mode == "q8_3")) comb_pending = true; // folded into the next gr_norm + else if (!(fold_comb && hc_mode != "3")) push(d, "shared", 0); // else folded into the next HC read + } + } + // ---- the final mixer and the head + gr_read("output_hc", true, XA, true); + gemv2(T("output.weight"), XA, lm, 248320, "lm head"); + { + ArgArgs a{248320, 120}; + mtl::Dispatch d1, d2; + d1.p = pam1; + d1.args = {mtl::Arg::B(lm), mtl::Arg::B(apv), mtl::Arg::B(api), mtl::Arg::V(a)}; + d1.grid[0] = 120; + d1.grid[1] = ntok; + d1.tg[0] = 256; + push(d1, "lm head", 0); + d2.p = pam2; + d2.args = {mtl::Arg::B(apv), mtl::Arg::B(api), mtl::Arg::B(tokout), mtl::Arg::V(a)}; + d2.grid[0] = ntok; + d2.tg[0] = 256; + push(d2, "lm head", 0); + } + // the producers' plane stride must be MAXT rows: gr_up's and swiglu's args carry n_tok for the stride, so the + // gr_up launches above use a.n_tok = ntok - patch them to MAXT (the kernels loop over GR_NTOK tokens anyway) + for (auto& d : ds) + if (d.p == pup || d.p == pupq || d.p == pnorm || d.p == pdown || d.p == pdownq) { + GrArgs a; + std::memcpy(&a, d.args.back().bytes.data(), sizeof a); + if (d.p == pup || d.p == pupq) a.n_tok = MAXT; + std::memcpy(d.args.back().bytes.data(), &a, sizeof a); + } + + double total_bytes = 0; + for (auto& kv : bytes_by) total_bytes += kv.second; + std::printf("%zu dispatches; weight + state bytes %.1f MB\n", ds.size(), total_bytes / 1e6); + auto run = [&] { return graph ? cx.run_graph(ds) : cx.run(ds); }; + run(); + const double tw = mtl::now_s(); + while (mtl::now_s() - tw < 1.0) run(); + std::vector t; + for (int i = 0; i < reps; ++i) t.push_back(run()); + std::printf("HC read: %s\n", hc_mode == "3" ? "BF16, 3 dispatches" : hc_mode == "2" ? "BF16, 2 dispatches" : hc_mode == "q8_3" ? "Q8_0, 3 dispatches" : "Q8_0, 2 dispatches"); + std::printf("encoding: %s\n", graph ? "concurrent encoder, barriers at dependencies" : "serial encoder"); + const double m = median(t); + std::printf("decode step (%d token(s)): GPU %.3f ms (min %.3f, max %.3f) -> %.1f GB/s effective; %.1f tok/s if " + "every token were a window of %d\n", ntok, m * 1e3, *std::min_element(t.begin(), t.end()) * 1e3, + *std::max_element(t.begin(), t.end()) * 1e3, total_bytes / m / 1e9, ntok / m, ntok); + for (auto& kv : bytes_by) std::printf(" %-10s %8.1f MB\n", kv.first.c_str(), kv.second / 1e6); + return 0; +} diff --git a/edge0-strata/metal/tests/wrap_split_test.cpp b/edge0-strata/metal/tests/wrap_split_test.cpp new file mode 100644 index 00000000..3d12aad5 --- /dev/null +++ b/edge0-strata/metal/tests/wrap_split_test.cpp @@ -0,0 +1,95 @@ +// metal/tests/wrap_split_test.cpp - Context::wrap(ptr, bytes, max_span) (#2338): a file longer than maxBufferLength is +// wrapped as overlapping no-copy buffers, and a dispatch that binds the buffer at ANY offset sees the bytes of that +// range, whole, through the chunk the offset selects (encode() and run_icb()). The file is sparse (no disk used +// beyond the few pages written). Exit 0 = every check passed. +// wrap_split_test [--factor 1.6] the file is factor x maxBufferLength +#include "mtl_rt.h" + +#include +#include +#include + +#include +#include +#include +#include +#include + +static int failures = 0; +static void check(bool ok, const std::string& what) { + std::printf("%s %s\n", ok ? "ok " : "FAIL", what.c_str()); + failures += !ok; +} + +int main(int argc, char** argv) { + double factor = 1.6; + for (int i = 1; i + 1 < argc; ++i) + if (!std::strcmp(argv[i], "--factor")) factor = std::atof(argv[i + 1]); + mtl::Context cx; + const size_t pg = 16384, maxlen = cx.max_buffer_length() / pg * pg; + const size_t bytes = ((size_t) (factor * (double) maxlen) + pg - 1) / pg * pg; + const size_t span = 256ull << 20; // the largest range bound at one offset: 256 MiB + std::printf("maxBufferLength %.2f GB, file %.2f GB, span %.0f MiB\n", maxlen / 1e9, bytes / 1e9, span / 1048576.0); + + char path[] = "/tmp/wrap_split_XXXXXX"; + const int fd = mkstemp(path); + if (fd < 0 || ftruncate(fd, (off_t) bytes) != 0) { std::printf("FAIL cannot make the sparse file\n"); return 1; } + unlink(path); + void* map = mmap(nullptr, bytes, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + if (map == MAP_FAILED) { std::printf("FAIL mmap\n"); return 1; } + + const int lib = cx.compile({"wrap_probe.metal"}); + const mtl::Pipeline pipe = cx.pipeline(lib, "probe_copy"); + mtl::Buffer W = cx.wrap(map, bytes, span); + const bool split = (bool) W.chunks; + if (factor > 1.0) { + check(split && W.chunks->size() >= 2, "the file is wrapped as " + std::to_string(W.chunks ? W.chunks->size() : 0) + " chunks"); + if (!split) return 1; + check(W.stride == maxlen - span, "stride = maxBufferLength - span"); + } + const size_t stride = split ? W.stride : bytes; // factor <= 1: one buffer, the same reads (a control) + + // starts of ranges bound at one offset: the file's start, either side of every chunk boundary, the last range + std::vector starts = {0, 4096, pg, bytes / 4 / pg * pg, bytes / 2 / pg * pg}; + for (size_t k = 1; k * stride < bytes; ++k) + for (long d : {-(long) pg, 0L, (long) pg, 3L * (long) pg}) { + const long s = (long) (k * stride) + d; + if (s > 0 && (size_t) s + span <= bytes) starts.push_back((size_t) s); + } + starts.push_back(bytes - span); + const uint32_t n = 4096, far = (uint32_t) (span / 4) - n; + mtl::Buffer dst = cx.alloc(2 * n * 4), par = cx.alloc(8); + auto* P = par.as(); + P[0] = n; P[1] = far; + int bad_enc = 0, bad_icb = 0; + for (size_t s : starts) { + s = s / 4 * 4; + auto* f = reinterpret_cast(static_cast(map) + s); + for (uint32_t i = 0; i < n; ++i) { f[i] = (float) (s / pg) * 10.0f + (float) i; f[far + i] = -(float) (s / pg) - 0.25f * (float) i; } + for (int mode = 0; mode < 2; ++mode) { + std::memset(dst.ptr, 0, dst.bytes); + mtl::Dispatch d; + d.p = &pipe; + d.args = {mtl::Arg::B(W, s), mtl::Arg::B(dst), mtl::Arg::B(par)}; + d.grid[0] = n / 64; d.tg[0] = 64; + if (mode == 0) cx.run({d}); else cx.run_icb({d}, 1, false); + const float* o = dst.as(); + int bad = 0; + for (uint32_t i = 0; i < n; ++i) + bad += o[i] != (float) (s / pg) * 10.0f + (float) i || o[n + i] != -(float) (s / pg) - 0.25f * (float) i; + (mode == 0 ? bad_enc : bad_icb) += bad != 0; + if (bad) std::printf(" offset %zu (%s): %d wrong values\n", s, mode == 0 ? "encode" : "icb", bad); + } + } + check(bad_enc == 0, "encode(): " + std::to_string(starts.size()) + " offsets, both ends of each range read right"); + // run_icb is a bench path (rt_bench only; the engine encodes with run/run_graph). On the M4 it reads wrong at some + // offsets of a single buffer too (--factor 0.9, no PR code involved), so it is reported, not counted. + std::printf("note run_icb(): %d of %zu offsets read wrong%s\n", bad_icb, starts.size(), split ? "" : " (single buffer)"); + cx.free(W); + cx.free(dst); + cx.free(par); + munmap(map, bytes); + close(fd); + std::printf(failures ? "FAILED\n" : "PASS\n"); + return failures ? 1 : 0; +} diff --git a/edge0-strata/metal/tools/chat_tokens.py b/edge0-strata/metal/tools/chat_tokens.py new file mode 100644 index 00000000..8bbf2c11 --- /dev/null +++ b/edge0-strata/metal/tools/chat_tokens.py @@ -0,0 +1,23 @@ +#!/usr/bin/env python3 +"""chat_tokens.py REPO TOKENIZER_DIR "question" : the token ids of a one-turn chat prompt (thinking off), comma +separated, for `strata generate --tokens` (the same template as the gates' rbench.py). +With --decode IDS instead of the question: the text of comma-separated ids.""" +import json, sys +from pathlib import Path + +repo, tdir = Path(sys.argv[1]), Path(sys.argv[2]) +sys.path.insert(0, str(repo / "tools")) +import strata_tokenizer as ST + +vocab = json.loads((tdir / "vocab.json").read_text(encoding="utf-8")) +toks = [None] * len(vocab) +for t, i in vocab.items(): + toks[i] = t +tok = ST.Tokenizer(toks, (tdir / "merges.txt").read_text(encoding="utf-8").split("\n"), + json.loads((tdir / "token_type.json").read_text())) +if sys.argv[3] == "--decode": + print(tok.decode([int(x) for x in sys.argv[4].replace(" ", "").split(",") if x])) +else: + ids = tok.encode("<|im_start|>user\n%s<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" % sys.argv[3], + parse_special=True) + print(",".join(str(i) for i in ids)) diff --git a/edge0-strata/metal/tools/check_ref.py b/edge0-strata/metal/tools/check_ref.py new file mode 100644 index 00000000..71da9290 --- /dev/null +++ b/edge0-strata/metal/tools/check_ref.py @@ -0,0 +1,34 @@ +"""metal/tools/check_ref.py - checks metal/src/ref_quant.cpp's dequantizers against gguf-py's (and Q2_0 against +tools/gguf_writer.py's codec) on the files tests/ref_dump wrote: + + python metal/tools/check_ref.py + +Prints the max |difference| per type; the reference passes when every one is 0 (bitwise: both decode to float32). +""" +import pathlib +import sys + +import numpy as np + +ROOT = pathlib.Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT / "tools")) +from gguf.quants import dequantize # noqa: E402 +from gguf.constants import GGMLQuantizationType as Q # noqa: E402 + +ROWS, COLS = 16, 2560 +d = pathlib.Path(sys.argv[1] if len(sys.argv) > 1 else ".") +bad = 0 +for name in ("Q2_0", "Q4_0", "Q5_0", "Q8_0", "Q3_K", "Q4_K", "Q5_K", "Q6_K", "IQ2_XXS", "IQ2_XS", "IQ2_S", "IQ3_XXS", + "IQ3_S", "IQ4_NL", "IQ4_XS", "BF16", "F16"): + raw = np.fromfile(d / f"{name}.bin", dtype=np.uint8) + ref = np.fromfile(d / f"{name}.f32", dtype=np.float32).reshape(ROWS, COLS) + if name == "Q2_0": + from gguf_writer import dequantize_q2_0 # noqa: E402 + got = np.asarray(dequantize_q2_0(raw.tobytes()), dtype=np.float32).reshape(ROWS, COLS) + else: + got = dequantize(raw.reshape(ROWS, -1), Q[name]).astype(np.float32).reshape(ROWS, COLS) + diff = float(np.max(np.abs(got - ref))) + scale = float(np.max(np.abs(got))) + print(f"{name:8s} max |ref_quant - gguf-py| = {diff:.3e} (max |w| {scale:.3e}) {'ok' if diff == 0 else 'DIFF'}") + bad += diff != 0 +sys.exit(1 if bad else 0) diff --git a/edge0-strata/metal/tools/gen_grids.sh b/edge0-strata/metal/tools/gen_grids.sh new file mode 100644 index 00000000..5fb5ff70 --- /dev/null +++ b/edge0-strata/metal/tools/gen_grids.sh @@ -0,0 +1,29 @@ +#!/bin/bash +# ggml declares its tables with GGML_TABLE_BEGIN(type, name, size) ... GGML_TABLE_END(); parse that form. +set -u +cd "$HOME/strata-mac/dsv41" || exit 1 +python3 - <<'PY' +import re, os +s = open("third_party/ggml/ggml-common.h").read() +out = ['''// metal/src/engine/msl/ggml_grids.h - ggml's i-quant grid tables, for the MSL expert kernels. +// +// Generated from third_party/ggml/ggml-common.h (pinned 3cf03257f): the CUDA and Metal kernels must use bit-identical tables, +// so they are copied from the same file rather than retyped. MSL cannot include a C header, so each array becomes a +// `constant` in the address space the kernels need - same values, same order. +#pragma once + +'''] +pat = re.compile(r"GGML_TABLE_BEGIN\(\s*(uint\d+_t)\s*,\s*(\w+)\s*,\s*([^)]+)\)(.*?)GGML_TABLE_END\(\)", re.S) +msl = {"uint8_t": "uchar", "uint16_t": "ushort", "uint32_t": "uint", "uint64_t": "ulong"} +got = [] +for ty, name, size, body in pat.findall(s): + if name not in ("ksigns_iq2xs", "kmask_iq2xs", "iq2xxs_grid", "iq2xs_grid", "iq3xxs_grid", "iq3s_grid", "iq2s_grid", "kgrid_q4nl"): + continue + vals = " ".join(body.split()) + out.append("constant %s %s[%s] = {%s};\n" % (msl[ty], name, size, vals)) + got.append("%s[%s]" % (name, size)) +open("metal/src/engine/msl/ggml_grids.h", "w").write("\n".join(out)) +print("emitted:", ", ".join(got) if got else "none") +PY +ls -l metal/src/engine/msl/ggml_grids.h | awk '{print $5, $9}' +grep -c "^constant" metal/src/engine/msl/ggml_grids.h diff --git a/edge0-strata/metal/tools/gen_stubs.py b/edge0-strata/metal/tools/gen_stubs.py new file mode 100644 index 00000000..f70ea34d --- /dev/null +++ b/edge0-strata/metal/tools/gen_stubs.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +"""gen_stubs.py LINK_LOG OUT.cpp CLANG_ARGS... : stubs for the kernel-library entry points the Metal engine leaves undefined. + +LINK_LOG is the output of a failed link of `strata` (ld64 prints each undefined symbol demangled). The declarations +come from the engine's own headers through clang's AST dump (-ast-dump-filter), so every stub has the header's exact +types; it is defined in the declaring namespace (free functions) or out of line (members, trailing return type), and +calls strata::metal_engine::missing("name") (metal/src/engine/missing.cpp). Entry points that have a Metal +implementation are defined in metal/src/engine/*.cpp and are therefore not undefined: they never get a stub. +""" +import os, re, subprocess, sys, tempfile + +log, out = sys.argv[1], sys.argv[2] +clang_args = sys.argv[3:] +root = None +for a in clang_args: + if a.startswith("-I") and a.endswith("/include") and os.path.isdir(os.path.join(a[2:], "strata")): + root = a[2:] +assert root, "need -I/include among the clang args" + +# ---- the undefined symbols (demangled), from ld64's report +undef = [] +for m in re.finditer(r'^\s+"(.+)", referenced from:', open(log, errors="replace").read(), re.M): + undef.append(m.group(1)) +undef = sorted(set(undef)) + +def split_top(s): + out, d, cur = [], 0, "" + for ch in s: + if ch in "(<[": d += 1 + if ch in ")>]": d -= 1 + if ch == "," and d == 0: + out.append(cur.strip()); cur = "" + else: + cur += ch + if cur.strip(): out.append(cur.strip()) + return out + +def parse_demangled(s): + """'ns::f(int, float*) const' -> ('ns::f', nparams)""" + s = re.sub(r"\s+const$", "", s) + d, i = 0, len(s) - 1 + while i >= 0: + if s[i] == ")": d += 1 + elif s[i] == "(": + d -= 1 + if d == 0: break + i -= 1 + name, params = s[:i], s[i + 1:-1] + p = split_top(params) + if p == ["void"]: p = [] + return name, len(p) + +wanted = {} +skipped = [] +for u in undef: + if "(" not in u or u.startswith("_") or "::" not in u: + skipped.append(u); continue + n, k = parse_demangled(u) + wanted.setdefault(n, set()).add(k) + +# ---- declarations from the headers +hdrs = [] +for d in ["kernels", "kernels/cpu", "prefill", "core", "platform", "ngram", "spec"]: + p = os.path.join(root, "strata", d) + if os.path.isdir(p): + hdrs += sorted("strata/%s/%s" % (d, f) for f in os.listdir(p) if f.endswith(".hpp")) +def compiles(h): + f = tempfile.NamedTemporaryFile("w", suffix=".cpp", delete=False) + f.write('#include "%s"\n' % h); f.close() + r = subprocess.run(["clang++", "-fsyntax-only"] + clang_args + [f.name], capture_output=True, text=True) + os.unlink(f.name) + return r.returncode == 0 +bad = [h for h in hdrs if not compiles(h)] +hdrs = [h for h in hdrs if h not in bad] +for h in bad: print("HEADER-SKIPPED", h) +tu = tempfile.NamedTemporaryFile("w", suffix=".cpp", delete=False) +for h in hdrs: tu.write('#include "%s"\n' % h) +tu.close() +decls = {} # qualname -> list of (kind, type_str, params[(name,type)], is_const, is_static, file) +cmd = ["clang++", "-fsyntax-only", "-Xclang", "-ast-dump", "-Xclang", "-ast-dump-filter=strata", "-fno-color-diagnostics"] + clang_args + [tu.name] +r = subprocess.run(cmd, capture_output=True, text=True) +if r.returncode != 0: sys.stderr.write(r.stderr[-4000:]) +# the AST dump is a tree ("|-", "`-", "| " prefixes, two characters per level): walk it, keeping the enclosing +# namespaces and records, and collect every function declaration that has no body +FN = ("FunctionDecl", "CXXMethodDecl", "CXXConstructorDecl", "CXXDestructorDecl") +stack = [] # (depth, kind, name) +cur = None # the function being read: [qual, kind, type, params, depth, has_body, rest] +def flush(): + global cur + if cur and not cur[5] and cur[0] in wanted: + q, kind, ty, params, d, hb, rest = cur + decls.setdefault(q, []).append((kind, ty, params, bool(re.search(r"\)\s*const\s*(noexcept)?\s*$", ty)), " static" in rest, "")) + cur = None +for line in r.stdout.splitlines(): + if line.startswith("Dumping "): flush(); stack = []; continue + m = re.match(r"^([| `-]*)(\w+) ", line) + if not m: continue + depth = len(m.group(1)) // 2 + kind = m.group(2) + if cur and depth <= cur[4]: flush() + while stack and stack[-1][0] >= depth: stack.pop() + if cur: + if kind == "CompoundStmt" and depth == cur[4] + 1: cur[5] = True + if kind == "ParmVarDecl" and depth == cur[4] + 1: + pm = re.search(r"(?:col|line):[\d:]+(?: used| referenced)?(?: (\w+))? '([^']*)'(?::'[^']*')?", line) + if pm: cur[3].append((pm.group(1) or "", pm.group(2))) + continue + if kind == "NamespaceDecl": + nm = re.search(r"> (?:prev 0x\w+ )?(?:line|col):[\d:]+ (\w+)", line) + stack.append((depth, kind, nm.group(1) if nm else "(anonymous)")) + elif kind in ("CXXRecordDecl", "ClassTemplateDecl"): + nm = re.search(r"(?:class|struct|union) (\w+)", line) + if " implicit " in line or not nm: stack.append((depth, kind, "(implicit)")) + else: stack.append((depth, kind, nm.group(1))) + elif kind in FN and " implicit " not in line: + fm = re.findall(r"(~?\w+) '([^']*)'(?::'[^']*')?", line) + if not fm: continue + name, ty = fm[-1] + rest = line[line.rfind("'") + 1:] + names = [s[2] for s in stack] + if any(n.startswith("(") for n in names) or any(s[1] == "ClassTemplateDecl" for s in stack): continue + cur = [("::".join(names + [name])), kind, ty, [], depth, False, rest] + elif kind in ("FunctionTemplateDecl",): + stack.append((depth, kind, "(template)")) +flush() + +def nx(ty): + return " noexcept" if re.search(r"\)\s*(const\s*)?noexcept\s*$", ty) else "" + +def ret_stmt(r): + if r == "void": return "" + if r.endswith("&"): # a reference: to a zero-initialised object of its own + return " static std::remove_cv_t> s{}; return s;" % r + return " return {};" + +def ret_of(ty): + d, i = 0, len(ty) - 1 + # the function type is "R (P...)" [ const]; find the '(' that opens the parameter list + s = re.sub(r"\)\s*(const)?\s*(noexcept)?\s*$", ")", ty) + i = len(s) - 1; d = 0 + while i >= 0: + if s[i] == ")": d += 1 + elif s[i] == "(": + d -= 1 + if d == 0: break + i -= 1 + return s[:i].strip() + +lines = ['// metal/src/engine/stubs.cpp - GENERATED by metal/tools/gen_stubs.py: do not edit.', + '// The kernel-library entry points with no Metal implementation yet: each stops the engine with its name', + '// (metal/src/engine/missing.cpp). Regenerate after porting one.', ''] +lines += ['#include "%s"' % h for h in hdrs] +lines += ['#include ', '', 'namespace strata::metal_engine { void missing(const char* name); }', ''] +done, unmatched = 0, [] +for qn in sorted(wanted): + ds = decls.get(qn, []) + ds = [d for d in ds if len(d[2]) in wanted[qn]] + if not ds: + unmatched.append(qn); continue + seen = set() + for kind, ty, params, is_const, is_static, loc in ds: + key = (ty, kind) + if key in seen: continue + seen.add(key) + ps = ", ".join("%s p%d" % (t, k) if "(*)" not in t else t.replace("(*)", "(*p%d)" % k) for k, (n, t) in enumerate(params)) + if kind == "FunctionDecl": + ns, fn = qn.rsplit("::", 1) + r = ret_of(ty) + body = 'strata::metal_engine::missing("%s");' % qn + body += ret_stmt(r) + lines.append("namespace %s { %s %s(%s)%s { %s } }" % (ns, r, fn, ps, nx(ty), body)) + elif kind == "CXXMethodDecl": + r = ret_of(ty) + body = 'strata::metal_engine::missing("%s");' % qn + body += ret_stmt(r) + lines.append("auto %s(%s)%s%s -> %s { %s }" % (qn, ps, " const" if is_const else "", nx(ty), r, body)) + elif kind in ("CXXConstructorDecl", "CXXDestructorDecl"): + lines.append('%s(%s) { strata::metal_engine::missing("%s"); }' % (qn, ps, qn)) + done += 1 +open(out, "w").write("\n".join(lines) + "\n") +print("undefined %d, stubs %d, unmatched %d, skipped %d" % (len(undef), done, len(unmatched), len(skipped))) +for u in unmatched: print("UNMATCHED", u) +for u in skipped: print("SKIPPED", u) +os.unlink(tu.name) diff --git a/edge0-strata/metal/tools/gen_tq_lut.py b/edge0-strata/metal/tools/gen_tq_lut.py new file mode 100644 index 00000000..c94d23c0 --- /dev/null +++ b/edge0-strata/metal/tools/gen_tq_lut.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +"""metal/tools/gen_tq_lut.py - writes the trellis codebook into metal/kernels/quant.metal. + +third_party/ggml/ggml-common-tq.h holds tq_lut_f16[4096]: 2048 codebook points of two fp16 values each. The Metal +kernels read a point as one 32-bit word (low half = the first value), so this script rewrites the table between the +"BEGIN tq_lut" and "END tq_lut" markers of quant.metal with tq_lut_u[2048] = lut[2i] | lut[2i+1] << 16. + + python metal/tools/gen_tq_lut.py rewrite quant.metal + python metal/tools/gen_tq_lut.py --check exit 1 when quant.metal differs from what would be written +""" +import pathlib +import re +import sys + +ROOT = pathlib.Path(__file__).resolve().parents[2] +SRC = ROOT / "third_party" / "ggml" / "ggml-common-tq.h" +DST = ROOT / "metal" / "kernels" / "quant.metal" +BEGIN, END = "// BEGIN tq_lut", "// END tq_lut" + + +def table(): + text = SRC.read_text() + m = re.search(r"tq_lut_f16\[4096\]\s*=\s*\{(.*?)\};", text, re.S) + vals = [int(v, 16) for v in re.findall(r"0x[0-9a-fA-F]+", m.group(1))] + assert len(vals) == 4096, len(vals) + return [vals[2 * i] | (vals[2 * i + 1] << 16) for i in range(2048)] + + +def block(): + words = table() + lines = [BEGIN + " (generated by metal/tools/gen_tq_lut.py from third_party/ggml/ggml-common-tq.h; do not edit)", + "constant uint tq_lut_u[2048] = {"] + for i in range(0, 2048, 8): + lines.append(" " + ", ".join("0x%08x" % w for w in words[i:i + 8]) + ",") + lines.append("};") + lines.append(END) + return "\n".join(lines) + + +def main(): + src = DST.read_text() + i, j = src.index(BEGIN), src.index(END) + len(END) + new = src[:i] + block() + src[j:] + if "--check" in sys.argv: + sys.exit(0 if new == src else 1) + DST.write_text(new) + + +if __name__ == "__main__": + main() diff --git a/edge0-strata/metal/tools/gguf_list.cpp b/edge0-strata/metal/tools/gguf_list.cpp new file mode 100644 index 00000000..1376170e --- /dev/null +++ b/edge0-strata/metal/tools/gguf_list.cpp @@ -0,0 +1,25 @@ +// metal/tools/gguf_list.cpp - tensor directory and scalar metadata of a GGUF (name type shape offset), for the +// Metal window's wiring. clang++ -std=c++20 -I include metal/tools/gguf_list.cpp -o gguf_list +#include "strata/artifact/gguf_reader.hpp" + +#include +#include + +int main(int argc, char** argv) { + if (argc < 2) { std::printf("usage: gguf_list [substring]\n"); return 2; } + strata::GgufFile g(argv[1]); + const std::string f = argc > 2 ? argv[2] : ""; + for (const auto& [k, v] : g.metadata()) { + if (!f.empty() && k.find(f) == std::string::npos) continue; + if (v.is_num()) std::printf("META %s = %g\n", k.c_str(), v.num()); + else if (v.type == strata::MetaType::STRING) std::printf("META %s = %s\n", k.c_str(), v.s.substr(0, 80).c_str()); + else std::printf("META %s = [%llu items]\n", k.c_str(), (unsigned long long) v.count); + } + for (const auto& t : g.tensors()) { + if (!f.empty() && t.name.find(f) == std::string::npos) continue; + std::printf("%s %s [", t.name.c_str(), t.type_name()); + for (size_t i = 0; i < t.shape.size(); ++i) std::printf(i ? ",%llu" : "%llu", (unsigned long long) t.shape[i]); + std::printf("] off %llu\n", (unsigned long long) t.offset); + } + return 0; +} diff --git a/edge0-strata/metal/tools/metal_repack.py b/edge0-strata/metal/tools/metal_repack.py new file mode 100644 index 00000000..02af864a --- /dev/null +++ b/edge0-strata/metal/tools/metal_repack.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +"""metal_repack.py: a multi-shard GGUF (Unsloth's UD-*_K_XL: 4-6 shards, shards cutting through layers) as ONE GGUF +for the Metal engine, which maps a single file into one GPU buffer (or streams its experts from it). + + python3 metal/tools/metal_repack.py --shard1 <...-00001-of-0000N.gguf> --out -metal.gguf [--verify] + +--verify: after writing, every tensor of the output is hashed (SHA-256) against its bytes in the source shard; any +difference stops with an error and removes the output. + +The output holds shard 1's metadata (without the split.* keys) and every tensor of every shard except the n-gram +table (per_layer_token_embd: it stays in its shard, read from the SSD by the host through --ple-gguf), byte for +byte. The dense tensors come first; then each layer's routed experts (gate, up, down) together. Every tensor starts +on a 16 KiB boundary (general.alignment = 16384: Apple Silicon's page). Prints the shard holding the n-gram table. +""" +from __future__ import annotations + +import argparse +import atexit +import hashlib +import os +import re +import struct +import sys +from pathlib import Path + +ALIGN = 16384 +SYNC_EVERY = 1 << 30 # the output is flushed to the disk every GiB: at most this much (plus a chunk) is ever dirty +CHUNK = 64 << 20 +LOCK_SUFFIX = ".repack.lock" # .repack.lock holds the repack's pid while it runs (setup_mac.py and the engine look) +# GGML type -> (bytes per block, values per block) +TYPE_SIZE = {0: (4, 1), 1: (2, 1), 30: (2, 1), 2: (18, 32), 3: (20, 32), 6: (22, 32), 7: (24, 32), 8: (34, 32), + 10: (84, 256), 11: (110, 256), 12: (144, 256), 13: (176, 256), 14: (210, 256), 16: (66, 256), + 17: (74, 256), 18: (98, 256), 19: (50, 256), 20: (18, 32), 21: (110, 256), 22: (82, 256), 23: (136, 256), + 42: (18, 64)} +PLE = "per_layer_token_embd.weight" + + +def nocache(f) -> None: + """Writes through this file skip the page cache (macOS F_NOCACHE; measured with mincore on a 4 GB write: 100% of the file cached + without it, 0% with it). A repack writes 100+ GB: through the cache that is 100+ GB of dirty pages competing with everything + else for memory (the 128 GB M5 Max keeps a 37 GB file-cache floor and compresses other programs' memory beyond it). READS are + not helped on this macOS (a cold 4 GB read was 98% cached afterwards with F_NOCACHE set); those pages are clean and evictable, + the harm of the old tool was the unbounded dirty output.""" + try: + import fcntl + fcntl.fcntl(f.fileno(), getattr(fcntl, "F_NOCACHE", 48), 1) + except (ImportError, OSError, ValueError): + pass + + +def drop_cache(f) -> None: + """Where F_NOCACHE does not exist (Linux, the BSDs): ask the kernel to drop this file's cached pages.""" + if hasattr(os, "posix_fadvise"): + try: + os.posix_fadvise(f.fileno(), 0, 0, os.POSIX_FADV_DONTNEED) + except OSError: + pass + + +def write_all(o, buf) -> None: + mv = memoryview(buf) + while len(mv): + mv = mv[o.write(mv):] + + +class Reader: + def __init__(self, path: Path): + self.path = path + self.f = open(path, "rb") + magic, self.version, nt, nkv = struct.unpack("<4sIQQ", self.f.read(24)) + if magic != b"GGUF": + sys.exit(f"{path}: not a GGUF file") + self.kv = [] # (key, raw bytes of the whole entry) + self.align = 32 + for _ in range(nkv): + start = self.f.tell() + key = self._str() + t = struct.unpack(" str: + n = struct.unpack(" list[Path]: + m = re.match(r"(.*)-00001-of-(\d{5})\.gguf$", shard1.name) + if not m: + return [shard1] + n = int(m.group(2)) + return [shard1.with_name(f"{m.group(1)}-{i:05d}-of-{n:05d}.gguf") for i in range(1, n + 1)] + + +def tensor_hash(f, pos: int, n: int) -> str: + h = hashlib.sha256() + f.seek(pos) + while n: + buf = f.read(min(n, CHUNK)) + if not buf: + break + h.update(buf) + n -= len(buf) + return h.hexdigest() + + +def verify(out: Path, data0: int, order: list) -> list[str]: + """The names of the tensors whose bytes in `out` differ from their source shard's (SHA-256 each).""" + bad = [] + with open(out, "rb") as o: + nocache(o) + r = Reader(out) # the output's own table, as a reader sees it + table = {t["name"]: t for t in r.tensors} + for i, t in enumerate(order): + mine = table.get(t["name"]) + if mine is None or mine["abs"] != data0 + t["new"] or mine["bytes"] != t["bytes"] or \ + tensor_hash(o, mine["abs"], t["bytes"]) != tensor_hash(t["src"].f, t["abs"], t["bytes"]): + bad.append(t["name"]) + if i % 50 == 0: + print(f"\r verify {i + 1} / {len(order)}", end="", flush=True) + r.f.close() + print() + return bad + + +def gguf_str(s: str) -> bytes: + b = s.encode("utf-8") + return struct.pack(" int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--shard1", required=True, type=Path) + ap.add_argument("--out", required=True, type=Path) + ap.add_argument("--verify", action="store_true", help="hash every tensor of the output against the source") + a = ap.parse_args() + shards = [Reader(p) for p in shard_paths(a.shard1)] + every = [t for s in shards for t in s.tensors] + ple = [t for t in every if t["name"] == PLE] + keep = [t for t in every if t["name"] != PLE] + + def layer_of(t): + m = re.match(r"blk\.(\d+)\.", t["name"]) + return int(m.group(1)) if m else -1 + dense = [t for t in keep if "_exps." not in t["name"]] + role = {"ffn_gate_exps": 0, "ffn_up_exps": 1, "ffn_down_exps": 2} + exps = sorted((t for t in keep if "_exps." in t["name"]), + key=lambda t: (layer_of(t), role.get(t["name"].split(".")[2], 9))) + order = dense + exps + kv = [raw for key, raw in shards[0].kv if not key.startswith("split.") and key != "general.alignment"] + kv.append(gguf_str("general.alignment") + struct.pack("= SYNC_EVERY: + os.fsync(o.fileno()) + drop_cache(o) + drop_cache(src) + since = 0 + done += t["bytes"] + if t is order[-1] or t["name"].endswith("ffn_down_exps.weight"): + print(f"\r {done / 2**30:6.1f} / {sum(x['bytes'] for x in order) / 2**30:.1f} GiB", end="", flush=True) + o.truncate(total) + os.fsync(o.fileno()) + print() + if a.verify: + bad = verify(tmp, data0, order) + if bad: + tmp.unlink() + sys.exit(f"verify FAILED: {len(bad)} tensors differ from the source, first {bad[0]} (output removed)") + print(f"verify: {len(order)} tensors, SHA-256 identical to the source shards") + os.replace(tmp, a.out) + print(f"{a.out}: {len(order)} tensors ({len(dense)} dense, {len(exps)} expert), {total / 2**30:.2f} GiB") + if ple: + print(f"n-gram table: {ple[0]['src'].path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/edge0-strata/metal/tools/stream_experts.py b/edge0-strata/metal/tools/stream_experts.py new file mode 100644 index 00000000..4eedafbe --- /dev/null +++ b/edge0-strata/metal/tools/stream_experts.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +"""stream_experts.py: the routed experts of a GGUF re-laid for SSD streaming - every expert's gate, up and down rows +next to each other, so the Metal engine reads a missed expert with one request instead of three (the M4 Mac mini's +SSD: 2.97 GB/s for one 1.75 MB read per expert against 2.30 GB/s for three 0.58 MB reads, 6 threads, F_NOCACHE). + + python3 metal/tools/stream_experts.py --gguf [--out .experts] [--verify] + +The file: a 64-byte header ("STRX", version 1, 48 layers, 512 experts, the source GGUF's size), then per layer its +offset in this file and the three slice sizes (4 x u64 each), then the experts, layer by layer, expert by expert: +gate slice, up slice, down slice - each byte for byte the GGUF's. --verify compares every expert with the GGUF. +The engine (metal/src/engine/stream.cpp) uses it when it sits beside the GGUF as .experts and its header +matches; otherwise it reads the GGUF as before. +""" +from __future__ import annotations + +import argparse +import hashlib +import os +import struct +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import metal_repack as R # noqa: E402 + +NL, NE = 48, 512 +ALIGN = 16384 + + +def plan(gguf: Path): + """The reader, each layer's (gate, up, down, slice sizes), each layer's offset in the file, the file's size, the + first data offset.""" + rd = R.Reader(gguf) + tab = {t["name"]: t for t in rd.tensors} + lay = [] + for l in range(NL): + g, u, d = (tab[f"blk.{l}.ffn_{k}_exps.weight"] for k in ("gate", "up", "down")) + lay.append((g, u, d, g["bytes"] // NE, u["bytes"] // NE, d["bytes"] // NE)) + head_len = 64 + NL * 32 + data0 = (head_len + ALIGN - 1) // ALIGN * ALIGN + offs, at = [], data0 + for g, u, d, gs, us, ds in lay: + offs.append(at) + at += NE * (gs + us + ds) + at = (at + ALIGN - 1) // ALIGN * ALIGN + return rd, lay, offs, at, data0 + + +def size_bytes(gguf: Path) -> int: + """The size the file for this GGUF will have (setup checks the free disk with it).""" + rd, _, _, at, _ = plan(gguf) + rd.f.close() + return at + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--gguf", required=True, type=Path) + ap.add_argument("--out", type=Path) + ap.add_argument("--verify", action="store_true") + a = ap.parse_args() + out = a.out or a.gguf.with_name(a.gguf.name + ".experts") + rd, lay, offs, at, data0 = plan(a.gguf) + head = struct.pack("<4sIIIQ", b"STRX", 1, NL, NE, os.path.getsize(a.gguf)).ljust(64, b"\0") + for (g, u, d, gs, us, ds), o in zip(lay, offs): + head += struct.pack(" 0) { @@ -2909,6 +2910,7 @@ bool Verifier::capture(int T, std::string& err) { } } } +#endif #if defined(STRATA_USE_HIP) // 143-amdk: the same listing on HIP (kernel names by hipKernelNameRefByPtr); STRATA_VERIFY_NODES=2 also lists the first nodes in graph order if (std::getenv("STRATA_VERIFY_NODES") != nullptr) { size_t nn = 0; diff --git a/edge0-strata/src/prefill/prefill.cpp b/edge0-strata/src/prefill/prefill.cpp index bf1db352..41f5cef8 100644 --- a/edge0-strata/src/prefill/prefill.cpp +++ b/edge0-strata/src/prefill/prefill.cpp @@ -5338,6 +5338,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri for (size_t q = 0; q <= nb2 && q < 17; ++q) std::fprintf(stderr, "latmmq-bh[%zu]=%d\n", q, bh[q]); } +#if !defined(STRATA_GPU_METAL) { // latmmq-dev: which device each buffer lives on, and the current one static bool devsaid = false; if (!devsaid && std::getenv("STRATA_LATENT_MMQ_TRACE") != nullptr) { @@ -5366,6 +5367,7 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri devof(m.lat_gdst)); } } +#endif auto run_p = [&](const mmq::Product& pr) { static const bool alt = std::getenv("STRATA_LATENT_MMQ_ALTCTX") != nullptr; static std::unique_ptr altctx; diff --git a/edge0-strata/src/program/generate.cpp b/edge0-strata/src/program/generate.cpp index 4a2b3fbe..97b0a7bd 100644 --- a/edge0-strata/src/program/generate.cpp +++ b/edge0-strata/src/program/generate.cpp @@ -236,6 +236,14 @@ struct LatentTab { if (d != nullptr) { strata::core::latent_shared_d() = d; strata::core::latent_shared_u() = u; return true; } const auto& lay0 = strata::kernels::cpu::expert_layout(); if (!(lay0.native && !lay0.fmt.empty() && lay0.fmt[0].n_embd != (int64_t) g.n_embd)) return true; +#if defined(STRATA_GPU_METAL) + const char* enabled = std::getenv("STRATA_METAL_LATENT"); + if (!enabled || std::strcmp(enabled, "1")) { + err = "enable the Metal latent path with STRATA_METAL_LATENT=1"; + return false; + } + return true; // The Metal window owns the latent files and GPU buffers. +#else latent = true; lat_H = (int64_t) g.n_embd; lat_ELL = lay0.fmt[0].n_embd; @@ -274,6 +282,7 @@ struct LatentTab { strata::core::latent_shared_d() = d; strata::core::latent_shared_u() = u; return true; +#endif } } lat_tab; // Windows' WDDM driver model: native Windows, or WSL2 (its GPU goes through /dev/dxg to the Windows driver). There, @@ -4151,6 +4160,11 @@ int main(int argc, char** argv) { const auto& lay = strata::kernels::cpu::expert_layout(); for (int64_t l = 0; lay.native && l < (int64_t) lay.fmt.size(); ++l) { const auto& f = lay.fmt[(size_t) l]; +#if defined(STRATA_GPU_METAL) + const char* latent = std::getenv("STRATA_METAL_LATENT"); + if (latent && !std::strcmp(latent, "1") && (f.n_embd == 1280 || f.n_embd == 1792) && f.n_ff == 640 && + f.gu_type == 11 && f.d_type == 20 && f.gu_colscale) continue; +#endif if (!strata::kernels::native_expert_supported(f.gu_type, f.d_type, f.n_embd, f.n_ff)) { std::fprintf(stderr, "strata generate: layer %lld's experts are %s/%s (ggml types %d/%d), which this " "engine has no GPU kernels for\n", (long long) l, diff --git a/edge0-strata/tools/configure_metal_latent.py b/edge0-strata/tools/configure_metal_latent.py new file mode 100644 index 00000000..d46242d0 --- /dev/null +++ b/edge0-strata/tools/configure_metal_latent.py @@ -0,0 +1,112 @@ +"""Write a local service configuration for the Q4_0 latent Metal backend.""" + +import argparse +import json +from pathlib import Path +import re + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--build-dir", type=Path, required=True) + parser.add_argument("--model-dir", type=Path, required=True) + parser.add_argument("--latent-dense", type=Path, required=True) + parser.add_argument("--tokenizer", type=Path, required=True) + parser.add_argument("--mtp-dir", type=Path, required=True) + parser.add_argument("--mtp-gguf", type=Path, required=True) + parser.add_argument("--expert-profile", type=Path) + parser.add_argument("--port", type=int, default=8127) + parser.add_argument("--output", type=Path, required=True) + options = parser.parse_args() + if not 1 <= options.port <= 65535: + parser.error("--port must be between 1 and 65535") + + source = Path(__file__).resolve().parents[1] + binary = options.build_dir.resolve() / "strata" + model = options.model_dir.resolve() + pack = model / "latent_a_pack" + dense = options.latent_dense.resolve() + mtp = options.mtp_dir.resolve() + mtp_gguf = options.mtp_gguf.resolve() + tokenizer = options.tokenizer.resolve() + files = [ + binary, pack / "native_experts.txt", pack / "experts.bin", + pack / "labels.bin", pack / "index.txt", pack / "dense.bin", + model / "model-metal-shell.gguf", model / "ple-fp8.gguf", + dense, mtp_gguf, mtp / "draft_vocab.bin", + ] + if options.expert_profile: + files.append(options.expert_profile.resolve()) + for path in files: + if not path.is_file(): + parser.error(f"missing file: {path}") + for path in (mtp, tokenizer): + if not path.is_dir(): + parser.error(f"missing directory: {path}") + + manifest = (pack / "native_experts.txt").read_text() + match = re.search(r"expert_embd (1280|1792)\b", manifest) + if not match: + parser.error("the expert manifest must declare latent width 1280 or 1792") + width = int(match[1]) + dense_bytes = 48 * 2 * 8 * 2560 * (width // 32) * 18 + if dense.stat().st_size != dense_bytes: + parser.error(f"Q4_0 D/U at width {width} requires {dense_bytes} bytes") + for layer in range(48): + scale = pack / "scales" / f"l{layer}.f32" + if not scale.is_file() or scale.stat().st_size != width * 4: + parser.error(f"expected {width} float32 scales in {scale}") + + args = [ + "--pack", str(pack), + "--native", str(model / "model-metal-shell.gguf"), + "--native-dense-gguf", str(model / "model-metal-shell.gguf"), + "--ple-gguf", str(model / "ple-fp8.gguf"), + "--ple-io", "direct", + "--expert-cache", "64", + "--mmap-experts", + "--prefill", "auto", + "--max-context", "32768", + "--spec", "4", + "--suffix-draft", "3", + "--mtp", str(mtp), + "--spec-min-p", "0.5", + ] + if options.expert_profile: + args += ["--expert-profile", str(options.expert_profile.resolve())] + env = { + "STRATA_METAL_LATENT": "1", + "STRATA_LATENT_DENSE": str(dense), + "STRATA_LATENT_DENSE_FMT": "4", + "STRATA_LATENT_A": str(pack / "experts.bin"), + "STRATA_LATENT_LABELS": str(pack / "labels.bin"), + "STRATA_HC_Q8": "0", + "STRATA_ROUTE_TAIL_SKIP": "0", + "STRATA_PREFILL_CPU_SHARE": "0", + "STRATA_METAL_LATENT_GROUP_U": "1", + "STRATA_METAL_LATENT_PROMPT": "1", + "STRATA_METAL_PROMPT_CHUNK": "8192", + "STRATA_METAL_LATENT_DECODE_TILES": "1", + "STRATA_METAL_MATH": "fast", + "STRATA_METAL_LATENT_PROMPT_QSA_MM": "1", + "STRATA_METAL_GDN_SCAN": "gdn_scanr_4_8", + "STRATA_METAL_MTP_GGUF": str(mtp_gguf), + "STRATA_METAL_DRAFT_VOCAB": str(mtp / "draft_vocab.bin"), + "STRATA_DECODE_TIMING": "1", + "STRATA_METAL_LATENT_VERIFY_FAST": "1", + "STRATA_METAL_LATENT_PROMPT_TUNE": "1", + } + output = options.output.resolve() + config = { + "exe": str(binary), "cwd": str(source), "args": args, "env": env, + "tokenizer": str(tokenizer), "model_name": f"strata-l{width}-q3k-metal", + "backend": "metal", "lib_dirs": [], "host": "127.0.0.1", + "port": options.port, "log": str(output.with_suffix(".engine.log")), + } + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(config, indent=2) + "\n") + print(f"Wrote {output}") + + +if __name__ == "__main__": + main() diff --git a/edge0-strata/tools/mtp_gguf.py b/edge0-strata/tools/mtp_gguf.py new file mode 100644 index 00000000..6e6d6727 --- /dev/null +++ b/edge0-strata/tools/mtp_gguf.py @@ -0,0 +1,146 @@ +"""Rebuild the Metal engine's MTP GGUF (mtp-q2_0.gguf) from the CUDA runtime directory (mtp/). + +The runtime directory is written by Strata's tools/mtp_rt.py from mtp-q2_0.gguf. This script inverts it: + + experts.bin Q2_0 blobs relaid out back to the fused gate_up_proj / down_proj tensors (byte-exact). + dense.bin bf16 entries are copied as BF16; f32 entries (RMSNorm weights stored as 1 + w) are written + as w; q8_0 entries are written as Q8_0 tensors. The Metal engine quantizes BF16 matrices to + Q8_0 when it loads the drafter and uses Q8_0 tensors as they are. + + python mtp_gguf.py --rt /mtp --out /mtp-q2_0.gguf +""" +from __future__ import annotations + +import argparse +import os +import struct +import sys +from pathlib import Path + +import numpy as np + +H, FF, NE = 2560, 640, 512 +Q2_BLOCK, Q2_BYTES = 64, 18 +GU_ROWS, GU_BLOCKS = 2 * FF, H // Q2_BLOCK # gate_up: 1280 rows of 2560 inputs +DN_ROWS, DN_BLOCKS = H, FF // Q2_BLOCK # down: 2560 rows of 640 inputs +BLOB = (GU_ROWS * GU_BLOCKS + DN_ROWS * DN_BLOCKS) * Q2_BYTES +ALIGN = 32 + +T_F32, T_Q8_0, T_BF16, T_Q2_0 = 0, 8, 30, 42 +GGUF_STRING, GGUF_UINT32 = 8, 4 + + +def experts(rt: Path) -> tuple[np.ndarray, np.ndarray]: + raw = np.fromfile(rt / "experts.bin", dtype=np.uint8) + if raw.size != NE * BLOB: + raise ValueError(f"experts.bin: {raw.size} bytes, expected {NE * BLOB}") + raw = raw.reshape(NE, BLOB) + gu_codes = GU_ROWS * GU_BLOCKS * 16 + dn_codes = DN_ROWS * DN_BLOCKS * 16 + gu_scales = GU_ROWS * GU_BLOCKS * 2 + o = 0 + parts = [] + for n in (gu_codes, dn_codes, gu_scales): + parts.append(raw[:, o:o + n]) + o += n + parts.append(raw[:, o:]) + gu = np.empty((NE, GU_ROWS, GU_BLOCKS, Q2_BYTES), dtype=np.uint8) + dn = np.empty((NE, DN_ROWS, DN_BLOCKS, Q2_BYTES), dtype=np.uint8) + inter = np.empty_like(gu) + inter[..., 2:] = parts[0].reshape(NE, GU_ROWS, GU_BLOCKS, 16) + inter[..., :2] = parts[2].reshape(NE, GU_ROWS, GU_BLOCKS, 2) + dn[..., 2:] = parts[1].reshape(NE, DN_ROWS, DN_BLOCKS, 16) + dn[..., :2] = parts[3].reshape(NE, DN_ROWS, DN_BLOCKS, 2) + gu[:, :FF] = inter[:, 0::2] # interleaved rows: 2r = gate r, 2r + 1 = up r + gu[:, FF:] = inter[:, 1::2] + return gu, dn + + +def full_name(short: str) -> str: + top = ("fc_embedding.", "fc_hidden.", "hyper_connection_mixer.", "pre_fc_norm_") + return "mtp." + short if short.startswith(top) else "mtp.layers.0." + short + + +def dense(rt: Path): + blob = (rt / "dense.bin").read_bytes() + for line in (rt / "dense.txt").read_text(encoding="utf-8").split("\n"): + if not line.strip(): + continue + short, kind, rows, cols, off, size = line.split() + rows, cols, off, size = int(rows), int(cols), int(off), int(size) + data = blob[off:off + size] + if kind == "f32": + plus_one = np.frombuffer(data, dtype="> 16) & 1)) & 0xFFFF0000).astype(np.uint32).view(np.float32) + w = np.where(snapped + np.float32(1.0) == plus_one, snapped, w).astype(np.float32) + if not np.array_equal(w + np.float32(1.0), plus_one): + raise ValueError(f"{short}: 1 + w does not round-trip") + yield full_name(short), T_F32, [cols], w.tobytes() + elif kind == "bf16": + if size != rows * cols * 2: + raise ValueError(f"{short}: {size} bytes for bf16 {rows}x{cols}") + yield full_name(short), T_BF16, [cols, rows], data + elif kind == "q8_0": + if cols % 32 or size != rows * cols // 32 * 34: + raise ValueError(f"{short}: {size} bytes for q8_0 {rows}x{cols}") + yield full_name(short), T_Q8_0, [cols, rows], data + else: + raise ValueError(f"{short}: unknown kind {kind}") + + +def gguf_string(s: str) -> bytes: + b = s.encode("utf-8") + return struct.pack(" None: + kv = [("general.architecture", "qwen4exp-mtp"), + ("strata.mtp.source", "Edge0 release mtp/ runtime files (tools/mtp_rt.py output), inverted"), + ("strata.mtp.expert_format", "q2_0")] + tensors = sorted(tensors, key=lambda t: t[0]) + head = bytearray(struct.pack("<4sIQQ", b"GGUF", 3, len(tensors), len(kv) + 1)) + head += gguf_string("general.alignment") + struct.pack(" int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--rt", required=True, help="the model's mtp/ directory (experts.bin, dense.bin, dense.txt)") + ap.add_argument("--out", required=True, help="output GGUF, e.g. /mtp-q2_0.gguf") + a = ap.parse_args() + rt = Path(a.rt) + gu, dn = experts(rt) + tensors = list(dense(rt)) + tensors.append(("mtp.layers.0.mlp.experts.gate_up_proj", T_Q2_0, [H, GU_ROWS, NE], gu.tobytes())) + tensors.append(("mtp.layers.0.mlp.experts.down_proj", T_Q2_0, [FF, DN_ROWS, NE], dn.tobytes())) + if len(tensors) != 31: + raise ValueError(f"{len(tensors)} tensors, expected 31") + write_gguf(Path(a.out), tensors) + print(f"{a.out}: {len(tensors)} tensors") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/edge0-strata/tools/verify_metal_latent.py b/edge0-strata/tools/verify_metal_latent.py new file mode 100644 index 00000000..44992197 --- /dev/null +++ b/edge0-strata/tools/verify_metal_latent.py @@ -0,0 +1,89 @@ +"""Check short answers and prompt chunks in a local latent Metal service.""" + +import argparse +import json +from pathlib import Path +import socket +import tempfile + +import strata_bench as bench + + +SHORT = [ + ("What is 17 times 23? Answer with the number only.", "391"), + ("What is the capital city of France? Answer with the city only.", "paris"), + ("How many days are in a leap year? Answer with the number only.", "366"), + ("What is 9 times 9? Answer with the number only.", "81"), + ("What is the capital of Japan? Answer with the city only.", "tokyo"), + ("Which planet is the largest in the solar system? Answer with the name only.", "jupiter"), + ("What is the chemical symbol for gold? Answer with the symbol only.", "au"), + ("What is 100 minus 37? Answer with the number only.", "63"), + ("On a clear day the sky is which colour? Answer with one word.", "blue"), + ("What is the capital of Italy? Answer with the city only.", "rome"), +] +LONG = [ + ("What is 17 times 23? Reply with just the number.", "391"), + ("How many days are in a leap year? Reply with just the number.", "366"), + ("Name the capital city of France. Reply with just the city.", "paris"), + ("What is 9 times 9? Reply with just the number.", "81"), +] +FILLER = ( + "The quick brown fox jumps over the lazy dog near the river bank while the old mill wheel turns slowly " + "and the miller counts sacks of grain in the fading afternoon light. " +) + + +def verify(base, output): + words = FILLER.split() + filler = " ".join((words * (8000 // len(words) + 1))[:8000]) + cases = [("short", question, question, expected) for question, expected in SHORT] + cases += [ + ("long", question, filler + "\n\nQuestion: " + question + "\nAnswer:", expected) + for question, expected in LONG + ] + results = [] + for kind, question, prompt, expected in cases: + reply = bench.chat(base, "", [{"role": "user", "content": prompt}], 24, think=False) + timing = reply.get("timings") or {} + tokens = (reply.get("usage") or {}).get("prompt_tokens") + answer = reply["text"].strip() + passed = not reply["err"] and expected in answer.lower() + if kind == "long": + passed = passed and tokens is not None and tokens > 8192 and timing.get("cache_n") == 0 + results.append({ + "kind": kind, "question": question, "expected": expected, "answer": answer, + "prompt_tokens": tokens, "reused_tokens": timing.get("cache_n"), + "passed": bool(passed), "error": reply["err"], + }) + print(f"{kind}: {tokens} prompt tokens, answer {answer!r}, passed={passed}", flush=True) + output.write_text(json.dumps(results, indent=2) + "\n") + return all(row["passed"] for row in results) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + options = parser.parse_args() + output = options.out.resolve() + output.parent.mkdir(parents=True, exist_ok=True) + config = json.loads(options.config.read_text()) + config["log"] = str(output.with_suffix(".engine.log")) + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + port = reserve.getsockname()[1] + with tempfile.TemporaryDirectory(prefix="strata-latent-") as temp: + config_path = Path(temp) / "config.json" + config_path.write_text(json.dumps(config)) + process = bench.start_server(config_path, port, None, output.with_suffix(".server.log")) + try: + base = bench.server_base(config.get("host"), port) + if not bench.wait_ready(base, "", process, 240): + raise RuntimeError("service did not load; check the engine and server logs") + return 0 if verify(base, output) else 1 + finally: + bench.stop_server(process) + + +if __name__ == "__main__": + raise SystemExit(main()) From 5c7f5c2378be9fef41b50a27fa309a0b922b165b Mon Sep 17 00:00:00 2001 From: shuyuzhang Date: Sat, 14 Nov 2026 15:05:07 +0800 Subject: [PATCH 2/4] edge0-infer: delegate Edge0 145B and Edge0 120B to edge0-strata on macOS Export of the development repository: the metal device and its option limits, the per-device template and the prepare step, contract 5, and the desktop app finding the edge0-infer in Contents/Resources of its macOS bundle. --- edge0-infer/README.md | 4 +- edge0-infer/README_es.md | 4 +- edge0-infer/README_fr.md | 4 +- edge0-infer/README_ja.md | 4 +- edge0-infer/README_zh.md | 4 +- edge0-infer/docs/architecture.md | 2 +- edge0-infer/docs/edge0-120b.md | 24 +- edge0-infer/docs/edge0-145b.md | 24 +- edge0-infer/docs/getting-started.md | 8 +- edge0-infer/docs/limitations.md | 2 +- edge0-infer/examples/desktop/README.md | 5 +- .../examples/desktop/src-tauri/src/service.rs | 44 +++- .../tests/integration/chat_strata_test.py | 39 ++- edge0-infer/tests/integration/fake_strata.py | 73 ++++++ .../tests/integration/serve_strata_test.py | 201 ++++++++++++++-- .../tests/integration/strata_support.py | 33 ++- .../tests/platform/strata/check_mtp_gguf.py | 114 +++++++++ .../tests/platform/strata/compare_outputs.py | 5 +- .../tests/unit/cli/strata_registry_test.cpp | 161 ++++++++++++- edge0-infer/tools/checks/check_layers.py | 8 +- edge0-infer/tools/cli/serve_delegate.h | 65 ++++- edge0-infer/tools/cli/serve_delegate_chat.cpp | 6 +- .../tools/cli/serve_delegate_commands.cpp | 28 +-- .../tools/cli/serve_delegate_process.cpp | 226 +++++++++++++++++- .../tools/cli/serve_delegate_proxy.cpp | 4 +- .../tools/cli/serve_delegate_registry.cpp | 169 +++++++++++-- .../tools/cli/serve_delegate_status.cpp | 2 +- 27 files changed, 1105 insertions(+), 158 deletions(-) create mode 100644 edge0-infer/tests/platform/strata/check_mtp_gguf.py diff --git a/edge0-infer/README.md b/edge0-infer/README.md index fdbaa9e3..aeac59c6 100644 --- a/edge0-infer/README.md +++ b/edge0-infer/README.md @@ -5,13 +5,13 @@ English | [Chinese](README_zh.md) | [Japanese](README_ja.md) | [Español](README `edge0-infer` runs mixture-of-experts (MoE) language models on your own hardware: a CLI, a local OpenAI- and Anthropic-compatible server and chat applications, with the same commands for every model. Each model runs on the engine built for it, and `edge0-infer` picks the engine from the model folder; there is nothing to choose: - **Edge0 35B** (Qwen3.6-35B-A3B with a LoRA and a trained pre-router) runs on the **native engine** in this directory, built on [GGML](https://github.com/ggml-org/ggml) and with a native C++ API, on macOS, iOS, Android, Linux and Windows. Expert weights are loaded from storage on demand into a fixed pool of slots, so a model larger than the device memory budget runs without loading all of its weights. On an Apple M5 Mac it decodes at about 25 tokens per second with a peak physical footprint of about 3.4 GiB; figures for every platform are in [Benchmarking](docs/benchmarking.md#reference-figures). -- **Edge0 145B** and **Edge0 120B** run on **edge0-strata**, a fork of [Strata](https://github.com/Niko1221/Strata) in `edge0-strata/` of the Edge0 repository, on an NVIDIA GPU on Linux or Windows; see [Edge0 145B](docs/edge0-145b.md) and [Edge0 120B](docs/edge0-120b.md). +- **Edge0 145B** and **Edge0 120B** run on **edge0-strata**, a fork of [Strata](https://github.com/Niko1221/Strata) in `edge0-strata/` of the Edge0 repository, on an NVIDIA GPU on Linux or Windows, and on Apple Silicon with Metal on macOS; see [Edge0 145B](docs/edge0-145b.md) and [Edge0 120B](docs/edge0-120b.md). ## Platforms | Platform | Edge0 35B (native engine) | Edge0 145B, Edge0 120B (edge0-strata) | | --- | --- | --- | -| macOS (Apple silicon) | Metal, CPU | — | +| macOS (Apple silicon) | Metal, CPU | Metal | | iOS | Metal | — | | Android | CPU | — | | Linux | CPU | NVIDIA GPU (CUDA 13) | diff --git a/edge0-infer/README_es.md b/edge0-infer/README_es.md index 8f7a1c59..d5eae734 100644 --- a/edge0-infer/README_es.md +++ b/edge0-infer/README_es.md @@ -5,7 +5,7 @@ `edge0-infer` ejecuta modelos de lenguaje de mezcla de expertos (MoE) en tu propio hardware: una CLI, un servidor local compatible con OpenAI y Anthropic y aplicaciones de chat, con los mismos comandos para todos los modelos. Cada modelo se ejecuta en el motor preparado para él, y `edge0-infer` elige el motor a partir de la carpeta del modelo; no hay nada que elegir: - **Edge0 35B** (Qwen3.6-35B-A3B con un LoRA y un pre-router entrenado) se ejecuta en el **motor nativo** de este directorio, construido sobre [GGML](https://github.com/ggml-org/ggml) y con una API nativa en C++, en macOS, iOS, Android, Linux y Windows. Los pesos de los expertos se cargan desde el almacenamiento bajo demanda en un conjunto fijo de ranuras, de modo que un modelo más grande que el presupuesto de memoria del dispositivo se ejecuta sin cargar todos sus pesos. En un Mac con Apple M5 decodifica a unos 25 tokens por segundo, con un pico de huella física de unos 3,4 GiB; las cifras de cada plataforma están en [Benchmarking](docs/benchmarking.md#reference-figures). -- **Edge0 145B** y **Edge0 120B** se ejecutan en **edge0-strata**, un fork de [Strata](https://github.com/Niko1221/Strata) en `edge0-strata/` del repositorio Edge0, en una GPU NVIDIA con Linux o Windows; consulta [Edge0 145B](docs/edge0-145b.md) y [Edge0 120B](docs/edge0-120b.md). +- **Edge0 145B** y **Edge0 120B** se ejecutan en **edge0-strata**, un fork de [Strata](https://github.com/Niko1221/Strata) en `edge0-strata/` del repositorio Edge0, en una GPU NVIDIA con Linux o Windows, y en Apple Silicon con Metal en macOS; consulta [Edge0 145B](docs/edge0-145b.md) y [Edge0 120B](docs/edge0-120b.md). > Esta página es una traducción del [README](README.md) en inglés; los documentos enlazados están en inglés. @@ -13,7 +13,7 @@ | Plataforma | Edge0 35B (motor nativo) | Edge0 145B, Edge0 120B (edge0-strata) | | --- | --- | --- | -| macOS (Apple silicon) | Metal, CPU | — | +| macOS (Apple silicon) | Metal, CPU | Metal | | iOS | Metal | — | | Android | CPU | — | | Linux | CPU | GPU NVIDIA (CUDA 13) | diff --git a/edge0-infer/README_fr.md b/edge0-infer/README_fr.md index 098152f0..fcac81d3 100644 --- a/edge0-infer/README_fr.md +++ b/edge0-infer/README_fr.md @@ -5,7 +5,7 @@ `edge0-infer` exécute des modèles de langage à mélange d'experts (MoE) sur votre propre matériel : une CLI, un serveur local compatible OpenAI et Anthropic et des applications de discussion, avec les mêmes commandes pour tous les modèles. Chaque modèle s'exécute sur le moteur conçu pour lui, et `edge0-infer` choisit le moteur d'après le dossier du modèle ; il n'y a rien à choisir : - **Edge0 35B** (Qwen3.6-35B-A3B avec un LoRA et un pre-router entraîné) s'exécute sur le **moteur natif** de ce répertoire, construit sur [GGML](https://github.com/ggml-org/ggml) et doté d'une API C++ native, sous macOS, iOS, Android, Linux et Windows. Les poids des experts sont chargés depuis le stockage à la demande dans un ensemble fixe d'emplacements, si bien qu'un modèle plus grand que le budget mémoire de l'appareil s'exécute sans charger tous ses poids. Sur un Mac Apple M5, il décode à environ 25 tokens par seconde, avec un pic d'empreinte physique d'environ 3,4 Gio ; les chiffres de chaque plateforme se trouvent dans [Benchmarking](docs/benchmarking.md#reference-figures). -- **Edge0 145B** et **Edge0 120B** s'exécutent sur **edge0-strata**, un fork de [Strata](https://github.com/Niko1221/Strata) dans `edge0-strata/` du dépôt Edge0, sur un GPU NVIDIA sous Linux ou Windows ; voir [Edge0 145B](docs/edge0-145b.md) et [Edge0 120B](docs/edge0-120b.md). +- **Edge0 145B** et **Edge0 120B** s'exécutent sur **edge0-strata**, un fork de [Strata](https://github.com/Niko1221/Strata) dans `edge0-strata/` du dépôt Edge0, sur un GPU NVIDIA sous Linux ou Windows, et sur Apple Silicon avec Metal sous macOS ; voir [Edge0 145B](docs/edge0-145b.md) et [Edge0 120B](docs/edge0-120b.md). > Cette page est une traduction du [README](README.md) en anglais ; les documents liés sont en anglais. @@ -13,7 +13,7 @@ | Plateforme | Edge0 35B (moteur natif) | Edge0 145B, Edge0 120B (edge0-strata) | | --- | --- | --- | -| macOS (Apple silicon) | Metal, CPU | — | +| macOS (Apple silicon) | Metal, CPU | Metal | | iOS | Metal | — | | Android | CPU | — | | Linux | CPU | GPU NVIDIA (CUDA 13) | diff --git a/edge0-infer/README_ja.md b/edge0-infer/README_ja.md index fa6b954d..74c6a137 100644 --- a/edge0-infer/README_ja.md +++ b/edge0-infer/README_ja.md @@ -5,7 +5,7 @@ `edge0-infer` は、混合エキスパート（MoE）言語モデルを手元のハードウェアで実行します。CLI、ローカルの OpenAI／Anthropic 互換サーバー、チャットアプリを備え、どのモデルも同じコマンドで使えます。各モデルはそれぞれ専用のエンジンで動作し、`edge0-infer` がモデルフォルダからエンジンを自動で選ぶため、選択する必要はありません。 - **Edge0 35B**（LoRA と学習済み pre-router を備えた Qwen3.6-35B-A3B）は、このディレクトリの**ネイティブエンジン**で動作します。ネイティブエンジンは [GGML](https://github.com/ggml-org/ggml) をベースにネイティブ C++ API を備え、macOS、iOS、Android、Linux、Windows で動作します。エキスパートの重みはストレージから必要に応じて固定数のスロットに読み込まれるため、デバイスのメモリ予算より大きなモデルでも、すべての重みを読み込まずに実行できます。Apple M5 Mac ではデコードが毎秒約 25 トークン、物理メモリフットプリントのピークが約 3.4 GiB です。各プラットフォームの数値は [Benchmarking](docs/benchmarking.md#reference-figures) を参照してください。 -- **Edge0 145B** と **Edge0 120B** は **edge0-strata** で動作します。[Strata](https://github.com/Niko1221/Strata) のフォークで、Edge0 リポジトリの `edge0-strata/` にあり、Linux または Windows の NVIDIA GPU で動作します。[Edge0 145B](docs/edge0-145b.md) と [Edge0 120B](docs/edge0-120b.md) を参照してください。 +- **Edge0 145B** と **Edge0 120B** は **edge0-strata** で動作します。[Strata](https://github.com/Niko1221/Strata) のフォークで、Edge0 リポジトリの `edge0-strata/` にあり、Linux または Windows の NVIDIA GPU、および macOS の Apple Silicon（Metal）で動作します。[Edge0 145B](docs/edge0-145b.md) と [Edge0 120B](docs/edge0-120b.md) を参照してください。 > このページは英語版 [README](README.md) の翻訳です。リンク先のドキュメントはすべて英語です。 @@ -13,7 +13,7 @@ | プラットフォーム | Edge0 35B（ネイティブエンジン） | Edge0 145B、Edge0 120B（edge0-strata） | | --- | --- | --- | -| macOS（Apple シリコン） | Metal、CPU | — | +| macOS（Apple シリコン） | Metal、CPU | Metal | | iOS | Metal | — | | Android | CPU | — | | Linux | CPU | NVIDIA GPU（CUDA 13） | diff --git a/edge0-infer/README_zh.md b/edge0-infer/README_zh.md index 638a2502..51801f29 100644 --- a/edge0-infer/README_zh.md +++ b/edge0-infer/README_zh.md @@ -5,7 +5,7 @@ `edge0-infer` 在你自己的硬件上运行混合专家（MoE）语言模型：命令行、本地 OpenAI／Anthropic 兼容服务与聊天应用，所有模型使用相同的命令。每个模型运行在为它准备的引擎上，`edge0-infer` 根据模型目录自动选择引擎，无需用户选择： - **Edge0 35B**（Qwen3.6-35B-A3B，带 LoRA 与训练好的 pre-router）运行在本目录的**原生引擎**上。原生引擎基于 [GGML](https://github.com/ggml-org/ggml)，提供原生 C++ API，支持 macOS、iOS、Android、Linux 和 Windows。专家权重从存储按需读入固定数量的专家槽，因此比设备内存预算更大的模型无需载入全部权重即可运行。在 Apple M5 Mac 上解码约每秒 25 个 token，物理内存占用峰值约 3.4 GiB；各平台的数据见 [Benchmarking](docs/benchmarking.md#reference-figures)。 -- **Edge0 145B** 和 **Edge0 120B** 运行在 **edge0-strata** 上。它是 [Strata](https://github.com/Niko1221/Strata) 的分支，位于 Edge0 仓库的 `edge0-strata/`，在 Linux 或 Windows 的 NVIDIA GPU 上运行；见 [Edge0 145B](docs/edge0-145b.md) 与 [Edge0 120B](docs/edge0-120b.md)。 +- **Edge0 145B** 和 **Edge0 120B** 运行在 **edge0-strata** 上。它是 [Strata](https://github.com/Niko1221/Strata) 的分支，位于 Edge0 仓库的 `edge0-strata/`，在 Linux 或 Windows 的 NVIDIA GPU 上运行，也在 macOS 的 Apple Silicon 上用 Metal 运行；见 [Edge0 145B](docs/edge0-145b.md) 与 [Edge0 120B](docs/edge0-120b.md)。 > 本页为英文版 [README](README.md) 的翻译；下面链接的文档均为英文。 @@ -13,7 +13,7 @@ | 平台 | Edge0 35B（原生引擎） | Edge0 145B、Edge0 120B（edge0-strata） | | --- | --- | --- | -| macOS（Apple Silicon） | Metal、CPU | — | +| macOS（Apple Silicon） | Metal、CPU | Metal | | iOS | Metal | — | | Android | CPU | — | | Linux | CPU | NVIDIA GPU（CUDA 13） | diff --git a/edge0-infer/docs/architecture.md b/edge0-infer/docs/architecture.md index 6f487e27..3d39f602 100644 --- a/edge0-infer/docs/architecture.md +++ b/edge0-infer/docs/architecture.md @@ -9,7 +9,7 @@ An overview of how `edge0-infer` is organized. It names the parts and how they w | Model | Engine | Hardware | How `edge0-infer` uses the engine | | --- | --- | --- | --- | | Edge0 35B | Native engine | GGML backends: Metal, CPU, Vulkan | A library in this directory, called in process; applications can also call it through the [native API](native-api.md) | -| Edge0 145B, Edge0 120B | edge0-strata | NVIDIA GPU with CUDA | A separate engine in `edge0-strata/` of the Edge0 repository, a fork of [Strata](https://github.com/Niko1221/Strata); `edge0-infer` starts it as a child process and talks to it over a loopback port | +| Edge0 145B, Edge0 120B | edge0-strata | NVIDIA GPU with CUDA, or Metal on macOS (Apple Silicon) | A separate engine in `edge0-strata/` of the Edge0 repository, a fork of [Strata](https://github.com/Niko1221/Strata); `edge0-infer` starts it as a child process and talks to it over a loopback port | `edge0-infer` picks the engine from the model folder given with `--model`; there is nothing to choose. A `model.json` with an `architecture` field is an asset manifest of the native engine. A `model.json` whose `model_name` names an Edge0 145B or Edge0 120B folder selects edge0-strata: `edge0-infer` starts the engine with that model's configuration template from `edge0-strata/edge0/models/`, waits until it is ready, and then forwards generation requests to it. `serve` keeps answering `/health`, `/v1/models` and `/status` itself, and `chat` keeps its own commands, so clients see the same interface for every model. The pages [Edge0 145B](edge0-145b.md) and [Edge0 120B](edge0-120b.md) describe this path; edge0-strata's own build and its launch contract with `edge0-infer` are in `edge0-strata/EDGE0.md`. diff --git a/edge0-infer/docs/edge0-120b.md b/edge0-infer/docs/edge0-120b.md index d2db844c..dd9e7dc1 100644 --- a/edge0-infer/docs/edge0-120b.md +++ b/edge0-infer/docs/edge0-120b.md @@ -1,15 +1,15 @@ # Edge0 120B -Edge0 120B is a mixture-of-experts model for a single NVIDIA GPU. It runs on **edge0-strata**, the engine built for it (see [Two layers](architecture.md#two-layers)): the `edge0-strata/` directory of the Edge0 repository, a fork of [Strata](https://github.com/Niko1221/Strata). `edge0-infer serve` and `edge0-infer chat` start edge0-strata for this model, wait until it is ready and talk to it on a loopback port, so clients use the same commands and endpoints as for every other model. [Edge0 145B](edge0-145b.md) runs the same way. +Edge0 120B is a mixture-of-experts model for one NVIDIA GPU on Linux or Windows, or for Apple Silicon on macOS. It runs on **edge0-strata**, the engine built for it (see [Two layers](architecture.md#two-layers)): the `edge0-strata/` directory of the Edge0 repository, a fork of [Strata](https://github.com/Niko1221/Strata). `edge0-infer serve` and `edge0-infer chat` start edge0-strata for this model, wait until it is ready and talk to it on a loopback port, so clients use the same commands and endpoints as for every other model. [Edge0 145B](edge0-145b.md) runs the same way. This page is the reference for running Edge0 120B with `edge0-infer`; [Getting started](getting-started.md#edge0-145b-and-edge0-120b-edge0-strata) walks through the steps. The model itself (architecture and measured speed) is described on the model page of Edge0 145B and Edge0 120B, `docs/models/edge0-145b.md` in the Edge0 repository. ## Requirements -- Linux or Windows with an NVIDIA GPU and CUDA 13. -- At least 16 GiB of GPU memory and 20 GiB of system memory. -- `edge0-infer` built as in [Build](getting-started.md#build) (the `cpu` preset on Linux) or [Windows](getting-started.md#windows); it needs no GPU backend of its own for these models. -- edge0-strata built as described in `edge0-strata/EDGE0.md` in the Edge0 repository. `edge0-infer` uses the engine `edge0-strata/build/strata` (`strata.exe` on Windows) and the Python environment `edge0-strata/.venv` it produces. +- Linux or Windows with an NVIDIA GPU and CUDA 13, or macOS on Apple Silicon with Metal. The macOS measurements were taken on an Apple M2 Ultra. +- On Linux and Windows, at least 16 GiB of GPU memory and 20 GiB of system memory. +- `edge0-infer` built as in [Build](getting-started.md#build) (the `cpu` preset on Linux, or the `cpu` or `mac-metal` preset on macOS) or [Windows](getting-started.md#windows); it needs no GPU backend of its own for these models. +- edge0-strata built as described in `edge0-strata/EDGE0.md` in the Edge0 repository. `edge0-infer` uses `edge0-strata/build/strata` on Linux, `edge0-strata/build/strata.exe` on Windows, or `edge0-strata/build-metal/strata` on macOS, and the Python environment `edge0-strata/.venv` that the build produces. ## Model directory @@ -53,7 +53,7 @@ build/cpu/bin/edge0-infer inspect --model "$EDGE0_120B" --json The report lists files with `exists` and `bytes`; `ok` is `true` (exit code 0) when the selected configuration is complete. The expert format comes from `model_name` and is not a model option; `--model-option experts=` is rejected. -D/U defaults to `q4`; add `--model-option du=q8` to check the other choice. +On Linux and Windows, D/U defaults to `q4`; add `--model-option du=q8` to check the other choice. On macOS the only accepted value is `q4`. Only Q3_K folders list or require `scales/` (in the report its `required` value is `experts=q3k`). When files are missing, standard output remains valid JSON and standard error lists the missing files. @@ -91,13 +91,13 @@ The commands are those of the native chat (`/file`, `"""`, `/think on|off`, `/ne | Option | Values | Default | | --- | --- | --- | -| `--device` | `cuda` or `cuda:` | `cuda:0` | -| `--model-option du=` | `q8` or `q4` | `q4` | +| `--device` | `metal` on macOS; `cuda` or `cuda:` on Linux and Windows | `metal` on macOS, `cuda:0` otherwise | +| `--model-option du=` | `q8` or `q4` | `q4` (`q4` on macOS) | | `--model-option ple-io=` | `direct` (read the embedding table from disk as needed) or `ram` (load it into memory; to keep it locked there, raise the locked-memory limit, e.g. `ulimit -l unlimited`) | `direct` | -| `--context` | tokens | 131072 | +| `--context` | tokens | 131072 (`32768` on macOS) | | `--strata-arg` | one word of edge0-strata's engine command line; repeat it for each word, e.g. `--strata-arg --pool-workers --strata-arg 24` | none | -The model settings (automatic expert cache and prefill, speculative decoding, and int8 KV cache) are in `edge0-strata/edge0/models/edge0-120b.json`. CPU worker count, PCIe mode and fraction, VRAM reserve, and KV residency use the engine's defaults unless given with `--strata-arg`. This is meant for tuning to one machine, for example: +On Linux and Windows the model settings (automatic expert cache and prefill, speculative decoding, and int8 KV cache) are in `edge0-strata/edge0/models/edge0-120b.json`. On macOS the same settings, with the limits below, are in `edge0-120b-metal.json`. CPU worker count, PCIe mode and fraction, VRAM reserve, and KV residency use the engine's defaults unless given with `--strata-arg`. This is meant for tuning to one machine, for example: ```sh build/cpu/bin/edge0-infer serve --model "$EDGE0_120B" \ @@ -106,7 +106,9 @@ build/cpu/bin/edge0-infer serve --model "$EDGE0_120B" \ --strata-arg --kv-resident --strata-arg 32768 ``` -The words are appended in the order given, after the template's arguments; write each option and its value as two words (not `--option=value`). An option the template already sets (such as `--kv` or `--pack`), `--max-context` (use `--context`) and an empty value are rejected before anything starts. Other words are not checked by `edge0-infer`: edge0-strata decides, and since the engine process's own output is not shown, a word it does not accept appears as a failed start (run edge0-strata directly to see why). The template requires launch contract version 4; an older engine template is rejected with both versions in the error. +The words are appended in the order given, after the template's arguments; write each option and its value as two words (not `--option=value`). An option the template already sets (such as `--kv` or `--pack`), `--max-context` (use `--context`) and an empty value are rejected before anything starts. Other words are not checked by `edge0-infer`: edge0-strata decides, and since the engine process's own output is not shown, a word it does not accept appears as a failed start (run edge0-strata directly to see why). On macOS, `--device` accepts only `metal`. An `fp4` folder is rejected. `du` accepts only `q4`, and `ple-io` accepts only `direct`. The Metal template `edge0-120b-metal.json` sets the context length to 32768. If `mtp-q2_0.gguf` is absent, the first start runs `tools/mtp_gguf.py` with edge0-strata's Python and writes that file into the model folder, which must be writable; a later start leaves the file in place. The file is not part of the released folder listing above. + +The template requires launch contract version 5; an older engine template is rejected with both versions in the error. ## CUDA libraries on Windows diff --git a/edge0-infer/docs/edge0-145b.md b/edge0-infer/docs/edge0-145b.md index d373190b..380a0a6b 100644 --- a/edge0-infer/docs/edge0-145b.md +++ b/edge0-infer/docs/edge0-145b.md @@ -1,15 +1,15 @@ # Edge0 145B -Edge0 145B is a mixture-of-experts model for a single NVIDIA GPU. It runs on **edge0-strata**, the engine built for it (see [Two layers](architecture.md#two-layers)): the `edge0-strata/` directory of the Edge0 repository, a fork of [Strata](https://github.com/Niko1221/Strata). `edge0-infer serve` and `edge0-infer chat` start edge0-strata for this model, wait until it is ready and talk to it on a loopback port, so clients use the same commands and endpoints as for every other model. [Edge0 120B](edge0-120b.md) runs the same way. +Edge0 145B is a mixture-of-experts model for one NVIDIA GPU on Linux or Windows, or for Apple Silicon on macOS. It runs on **edge0-strata**, the engine built for it (see [Two layers](architecture.md#two-layers)): the `edge0-strata/` directory of the Edge0 repository, a fork of [Strata](https://github.com/Niko1221/Strata). `edge0-infer serve` and `edge0-infer chat` start edge0-strata for this model, wait until it is ready and talk to it on a loopback port, so clients use the same commands and endpoints as for every other model. [Edge0 120B](edge0-120b.md) runs the same way. This page is the reference for running Edge0 145B with `edge0-infer`; [Getting started](getting-started.md#edge0-145b-and-edge0-120b-edge0-strata) walks through the steps. The model itself (architecture and measured speed) is described on the model page of Edge0 145B and Edge0 120B, `docs/models/edge0-145b.md` in the Edge0 repository. ## Requirements -- Linux or Windows with an NVIDIA GPU and CUDA 13. -- At least 16 GiB of GPU memory and 20 GiB of system memory. -- `edge0-infer` built as in [Build](getting-started.md#build) (the `cpu` preset on Linux) or [Windows](getting-started.md#windows); it needs no GPU backend of its own for these models. -- edge0-strata built as described in `edge0-strata/EDGE0.md` in the Edge0 repository. `edge0-infer` uses the engine `edge0-strata/build/strata` (`strata.exe` on Windows) and the Python environment `edge0-strata/.venv` it produces. +- Linux or Windows with an NVIDIA GPU and CUDA 13, or macOS on Apple Silicon with Metal. The macOS measurements were taken on an Apple M2 Ultra. +- On Linux and Windows, at least 16 GiB of GPU memory and 20 GiB of system memory. +- `edge0-infer` built as in [Build](getting-started.md#build) (the `cpu` preset on Linux, or the `cpu` or `mac-metal` preset on macOS) or [Windows](getting-started.md#windows); it needs no GPU backend of its own for these models. +- edge0-strata built as described in `edge0-strata/EDGE0.md` in the Edge0 repository. `edge0-infer` uses `edge0-strata/build/strata` on Linux, `edge0-strata/build/strata.exe` on Windows, or `edge0-strata/build-metal/strata` on macOS, and the Python environment `edge0-strata/.venv` that the build produces. ## Model directory @@ -53,7 +53,7 @@ build/cpu/bin/edge0-infer inspect --model "$EDGE0_145B" --json The report lists files with `exists` and `bytes`; `ok` is `true` (exit code 0) when the selected configuration is complete. The expert format comes from `model_name` and is not a model option; `--model-option experts=` is rejected. -D/U defaults to `q8`; add `--model-option du=q4` to check the other choice. +On Linux and Windows, D/U defaults to `q8`; add `--model-option du=q4` to check the other choice. On macOS the only accepted value is `q4`. Only Q3_K folders list or require `scales/` (in the report its `required` value is `experts=q3k`). When files are missing, standard output remains valid JSON and standard error lists the missing files. @@ -91,13 +91,13 @@ The commands are those of the native chat (`/file`, `"""`, `/think on|off`, `/ne | Option | Values | Default | | --- | --- | --- | -| `--device` | `cuda` or `cuda:` | `cuda:0` | -| `--model-option du=` | `q8` or `q4` | `q8` | +| `--device` | `metal` on macOS; `cuda` or `cuda:` on Linux and Windows | `metal` on macOS, `cuda:0` otherwise | +| `--model-option du=` | `q8` or `q4` | `q8` (`q4` on macOS) | | `--model-option ple-io=` | `direct` (read the embedding table from disk as needed) or `ram` (load it into memory; to keep it locked there, raise the locked-memory limit, e.g. `ulimit -l unlimited`) | `direct` | -| `--context` | tokens | 131072 | +| `--context` | tokens | 131072 (`32768` on macOS) | | `--strata-arg` | one word of edge0-strata's engine command line; repeat it for each word, e.g. `--strata-arg --pool-workers --strata-arg 24` | none | -The model settings (automatic expert cache and prefill, speculative decoding, and int8 KV cache) are in `edge0-strata/edge0/models/edge0-145b.json`. CPU worker count, PCIe mode and fraction, VRAM reserve, and KV residency use the engine's defaults unless given with `--strata-arg`. This is meant for tuning to one machine, for example: +On Linux and Windows the model settings (automatic expert cache and prefill, speculative decoding, and int8 KV cache) are in `edge0-strata/edge0/models/edge0-145b.json`. On macOS the same settings, with the limits below, are in `edge0-145b-metal.json`. CPU worker count, PCIe mode and fraction, VRAM reserve, and KV residency use the engine's defaults unless given with `--strata-arg`. This is meant for tuning to one machine, for example: ```sh build/cpu/bin/edge0-infer serve --model "$EDGE0_145B" \ @@ -106,7 +106,9 @@ build/cpu/bin/edge0-infer serve --model "$EDGE0_145B" \ --strata-arg --kv-resident --strata-arg 32768 ``` -The words are appended in the order given, after the template's arguments; write each option and its value as two words (not `--option=value`). An option the template already sets (such as `--kv` or `--pack`), `--max-context` (use `--context`) and an empty value are rejected before anything starts. Other words are not checked by `edge0-infer`: edge0-strata decides, and since the engine process's own output is not shown, a word it does not accept appears as a failed start (run edge0-strata directly to see why). The template requires launch contract version 4; an older engine template is rejected with both versions in the error. +The words are appended in the order given, after the template's arguments; write each option and its value as two words (not `--option=value`). An option the template already sets (such as `--kv` or `--pack`), `--max-context` (use `--context`) and an empty value are rejected before anything starts. Other words are not checked by `edge0-infer`: edge0-strata decides, and since the engine process's own output is not shown, a word it does not accept appears as a failed start (run edge0-strata directly to see why). On macOS, `--device` accepts only `metal`. An `fp4` folder is rejected. `du` accepts only `q4`, and `ple-io` accepts only `direct`. The Metal template `edge0-145b-metal.json` sets the context length to 32768. If `mtp-q2_0.gguf` is absent, the first start runs `tools/mtp_gguf.py` with edge0-strata's Python and writes that file into the model folder, which must be writable; a later start leaves the file in place. The file is not part of the released folder listing above. + +The template requires launch contract version 5; an older engine template is rejected with both versions in the error. ## CUDA libraries on Windows diff --git a/edge0-infer/docs/getting-started.md b/edge0-infer/docs/getting-started.md index 37f19b22..e6519563 100644 --- a/edge0-infer/docs/getting-started.md +++ b/edge0-infer/docs/getting-started.md @@ -9,11 +9,11 @@ From a fresh clone to a chat and a local OpenAI- and Anthropic-compatible server | Model | Engine | Hardware | Where the engine is | | --- | --- | --- | --- | | Edge0 35B | Native engine | Metal, CPU or Vulkan on macOS, iOS, Android, Linux and Windows | this directory, also usable in process through the [native API](native-api.md) | -| Edge0 145B, Edge0 120B | edge0-strata | NVIDIA GPU (CUDA 13) on Linux and Windows | `edge0-strata/` in the Edge0 repository; `edge0-infer` starts it as a separate process | +| Edge0 145B, Edge0 120B | edge0-strata | NVIDIA GPU (CUDA 13) on Linux and Windows, or Metal on macOS (Apple Silicon) | `edge0-strata/` in the Edge0 repository; `edge0-infer` starts it as a separate process | [Architecture](architecture.md) describes how the two layers work together. -Edge0 35B is verified on macOS on Apple silicon (Metal and CPU), on Android (CPU, PKR110), on Windows (CPU on an Intel Core i5-12400; Vulkan on an NVIDIA GeForce RTX 4070 SUPER) and on iOS (Metal, iPhone 17 Pro Max); measured speed and memory are under [Benchmarking](benchmarking.md#reference-figures). The measured speed of Edge0 145B and Edge0 120B is on their model page, `docs/models/edge0-145b.md` in the Edge0 repository. +Edge0 35B is verified on macOS on Apple silicon (Metal and CPU), on Android (CPU, PKR110), on Windows (CPU on an Intel Core i5-12400; Vulkan on an NVIDIA GeForce RTX 4070 SUPER) and on iOS (Metal, iPhone 17 Pro Max); measured speed and memory are under [Benchmarking](benchmarking.md#reference-figures). The measured speed of Edge0 145B and Edge0 120B is on their model page, `docs/models/edge0-145b.md` in the Edge0 repository. The macOS measurements were taken on an Apple M2 Ultra. ## Requirements @@ -86,9 +86,9 @@ The [asset tools guide](../tools/assets/README.md) describes the asset directory ### Edge0 145B and Edge0 120B (edge0-strata) -These models come as released model folders, one model and one expert format per folder (for example `Edge0-145B-q3k`); the asset tools above do not make them. They need an NVIDIA GPU with CUDA 13 on Linux or Windows. +These models come as released model folders, one model and one expert format per folder (for example `Edge0-145B-q3k`); the asset tools above do not make them. They need an NVIDIA GPU with CUDA 13 on Linux or Windows, or Apple Silicon with Metal on macOS. On macOS only the `q3k` folder is accepted, with `du=q4` and `ple-io=direct`. The first start writes `mtp-q2_0.gguf` into the model folder, so that folder must be writable. -1. **Build edge0-strata**, the `edge0-strata/` directory beside `edge0-infer/` in the Edge0 repository, as its `EDGE0.md` describes (sections "Build (Linux, NVIDIA)" and "Build (Windows, NVIDIA)"). `edge0-infer` finds it there and uses the engine and the Python environment that the build produces. +1. **Build edge0-strata**, the `edge0-strata/` directory beside `edge0-infer/` in the Edge0 repository, as its `EDGE0.md` describes (sections "Build (Linux, NVIDIA)", "Build (Windows, NVIDIA)" and "Build (macOS, Apple Silicon)"). `edge0-infer` finds it there and uses the engine and the Python environment that the build produces. On macOS the engine is `build-metal/strata`. 2. **Check the model folder**, the one that contains `model.json`: diff --git a/edge0-infer/docs/limitations.md b/edge0-infer/docs/limitations.md index d90a9aa8..a7d3bcf4 100644 --- a/edge0-infer/docs/limitations.md +++ b/edge0-infer/docs/limitations.md @@ -10,4 +10,4 @@ Behavior of the current release that users and integrators should be aware of. ## Edge0 145B and Edge0 120B run only through serve and chat -Edge0 145B and Edge0 120B run on the edge0-strata engine, not on the native engine (see [Two layers](architecture.md#two-layers)): only through `edge0-infer serve` and `edge0-infer chat` (and `inspect`), and not through the native API. Their server queues concurrent requests instead of answering 429, and part of `/status` is null for them (no memory breakdown or expert statistics; progress of a request in flight is counted from the forwarded text fragments). See Known limitations for [Edge0 145B](edge0-145b.md#known-limitations) and [Edge0 120B](edge0-120b.md#known-limitations). +Edge0 145B and Edge0 120B run on the edge0-strata engine, not on the native engine (see [Two layers](architecture.md#two-layers)): only through `edge0-infer serve` and `edge0-infer chat` (and `inspect`), and not through the native API. On macOS they use Metal on Apple Silicon: only a `q3k` folder, `du=q4` and `ple-io=direct`, with a context length of 32768, and the first start writes `mtp-q2_0.gguf` into the model folder. Their server queues concurrent requests instead of answering 429, and part of `/status` is null for them (no memory breakdown or expert statistics; progress of a request in flight is counted from the forwarded text fragments). See Known limitations for [Edge0 145B](edge0-145b.md#known-limitations) and [Edge0 120B](edge0-120b.md#known-limitations). diff --git a/edge0-infer/examples/desktop/README.md b/edge0-infer/examples/desktop/README.md index a300a3d8..717dbdf9 100644 --- a/edge0-infer/examples/desktop/README.md +++ b/edge0-infer/examples/desktop/README.md @@ -146,6 +146,9 @@ Build `edge0-infer` using the [build guide](../../docs/getting-started.md). On Choose a model, expand Other location… to enter the full executable path and one asset directory per line. This section starts collapsed, including on first launch before an executable has been configured. +With no executable configured, the app uses the `edge0-infer` that a release +package ships with it: beside its own executable, or in `Contents/Resources` of +the macOS app bundle. If no executable is configured, or the configured file has been moved or deleted, loading or saving reports which one it is and points to Settings > Engine > Executable and Other location… to set it again. @@ -165,7 +168,7 @@ The default device is `metal` on macOS and `cpu` on Windows. Settings also allow chosen backend. Model loading can take several minutes. Cancel stops a pending load. Changing model or saving settings ends the old service and clears history. A model that `inspect --json` reports as `delegated` is started without the device -settings: it runs on the GPU its external engine uses. Its `edge0-infer` finds edge0-strata through +settings. On macOS it runs on Metal through edge0-strata; on Linux and Windows it runs on the NVIDIA GPU that engine uses. Its `edge0-infer` finds edge0-strata through `EDGE0_STRATA_DIR` in the app's environment or an `edge0-strata/` directory beside an ancestor of the selected executable. diff --git a/edge0-infer/examples/desktop/src-tauri/src/service.rs b/edge0-infer/examples/desktop/src-tauri/src/service.rs index 4708882b..411af2fa 100644 --- a/edge0-infer/examples/desktop/src-tauri/src/service.rs +++ b/edge0-infer/examples/desktop/src-tauri/src/service.rs @@ -296,7 +296,7 @@ impl Manager { } else { Settings::default() }; - // No executable chosen yet: use the edge0-infer that the release package puts beside this app. + // No executable chosen yet: use the edge0-infer that the release package ships with this app. if settings.executable.trim().is_empty() { if let Some(found) = std::env::current_exe() .ok() @@ -829,11 +829,21 @@ impl Drop for Manager { let _ = self.shutdown(); } } -/// The edge0-infer in `app_dir` (the directory of this app's own executable), as the release package lays -/// them out side by side; None when it is not there (a development build, an app bundle). +/// The edge0-infer that a release package ships with this app, given `app_dir`, the directory of this app's +/// own executable: beside it (the Windows package), or in the bundle's `Contents/Resources` when the app runs +/// from `Contents/MacOS` of a macOS app bundle. None when neither is there (a development build). pub(crate) fn bundled_executable(app_dir: &Path) -> Option { - let path = app_dir.join(format!("edge0-infer{}", std::env::consts::EXE_SUFFIX)); - path.is_file().then(|| path.to_string_lossy().into_owned()) + let name = format!("edge0-infer{}", std::env::consts::EXE_SUFFIX); + let mut candidates = vec![app_dir.join(&name)]; + if app_dir.file_name().is_some_and(|n| n == "MacOS") { + if let Some(contents) = app_dir.parent().filter(|c| c.file_name().is_some_and(|n| n == "Contents")) { + candidates.push(contents.join("Resources").join(&name)); + } + } + candidates + .into_iter() + .find(|path| path.is_file()) + .map(|path| path.to_string_lossy().into_owned()) } #[cfg(test)] @@ -864,6 +874,30 @@ mod tests { fs::remove_dir_all(dir).unwrap(); } #[test] + fn finds_the_bundled_executable_in_the_app_bundle_resources() { + let root = directory(); + let name = format!("edge0-infer{}", std::env::consts::EXE_SUFFIX); + let macos = root.join("Edge0 Chat.app").join("Contents").join("MacOS"); + let resources = macos.parent().unwrap().join("Resources"); + fs::create_dir_all(&macos).unwrap(); + fs::create_dir_all(&resources).unwrap(); + assert_eq!(bundled_executable(&macos), None); + let inside = resources.join(&name); + fs::write(&inside, b"").unwrap(); + assert_eq!(bundled_executable(&macos), Some(inside.to_string_lossy().into_owned())); + // The one beside the executable comes first. + let beside = macos.join(&name); + fs::write(&beside, b"").unwrap(); + assert_eq!(bundled_executable(&macos), Some(beside.to_string_lossy().into_owned())); + // Resources is only looked at from Contents/MacOS. + let other = root.join("bin"); + fs::create_dir_all(&other).unwrap(); + fs::create_dir_all(root.join("Resources")).unwrap(); + fs::write(root.join("Resources").join(&name), b"").unwrap(); + assert_eq!(bundled_executable(&other), None); + fs::remove_dir_all(root).unwrap(); + } + #[test] fn memory_is_steady_only_after_a_full_window_within_tolerance() { let gb = |v: f64| (v * GIB as f64) as u64; let mut committed: VecDeque = (0..STEADY_SAMPLES - 1).map(|_| gb(90.0)).collect(); diff --git a/edge0-infer/tests/integration/chat_strata_test.py b/edge0-infer/tests/integration/chat_strata_test.py index ea01a769..5c9cffd8 100644 --- a/edge0-infer/tests/integration/chat_strata_test.py +++ b/edge0-infer/tests/integration/chat_strata_test.py @@ -180,8 +180,12 @@ def failures(cli, root, engine, model): kept_log(record, result.stderr) assert_no_leftovers(record) record = root / 'rejected' - result, _ = chat(cli, model, engine, record, 'hi\n', ['--device', 'metal']) - assert result.returncode == 2 and 'CUDA GPU' in result.stderr and not record.exists() + if sys.platform == 'darwin': + result, _ = chat(cli, model, engine, record, 'hi\n', ['--device', 'cuda']) + assert result.returncode == 2 and 'must be metal' in result.stderr and not record.exists() + else: + result, _ = chat(cli, model, engine, record, 'hi\n', ['--device', 'metal']) + assert result.returncode == 2 and 'CUDA GPU' in result.stderr and not record.exists() result, _ = chat(cli, model, engine, record, 'hi\n', ['--temperature', '0']) assert result.returncode == 2 and '--temperature' in result.stderr and not record.exists() print('PASS failures: start failure and engine exit exit 1 with the log tail and a kept log, rejected options start nothing') @@ -220,14 +224,35 @@ def main(): cancel(args.cli, root, engine, model) failures(args.cli, root, engine, model) passthrough(args.cli, root, engine, model) + if sys.platform == 'darwin': + assert (model / 'mtp-q2_0.gguf').read_bytes() == b'mtp-q2_0' + failed_model = make_model(root / 'prepare-fail-model') + failed_record = root / 'prepare-fail' + failed_env = environment(failed_record) + failed_env['FAKE_MTP_MODE'] = 'fail' + failed = subprocess.run([str(args.cli), 'chat', '--model', str(failed_model), *strata_args(engine)], + input='hi\n', capture_output=True, encoding='utf-8', env=failed_env, timeout=120) + assert failed.returncode == 1 and 'preparation failed' in failed.stderr, failed.stderr + assert 'log kept at' not in failed.stderr + assert not (failed_model / 'mtp-q2_0.gguf').exists() + assert not list(temp_dir(failed_record).iterdir()), list(temp_dir(failed_record).iterdir()) + assert len((model / 'prepare-calls.jsonl').read_text().splitlines()) == 1 + again, _ = chat(args.cli, model, engine, root / 'chat-again', 'hi\n') + assert again.returncode == 0 and 'preparing' not in again.stderr, again.stderr + assert len((model / 'prepare-calls.jsonl').read_text().splitlines()) == 1 + assert_no_leftovers(root / 'chat-again') fp4 = make_model(root / 'fp4', model_name='edge0-120b-fp4') record = root / 'fp4-chat' result, bodies = chat(args.cli, fp4, engine, record, 'hi\n') - assert result.returncode == 0 and bodies[0]['model'] == 'edge0_120b', result.stderr - config = json.loads((record / 'config.json').read_text()) - assert config['env']['STRATA_LAT_FUSED_Q3K'] == '0' and not (fp4 / 'scales').exists() - assert config['args'][1] == str(fp4) - assert_no_leftovers(record) + if sys.platform == 'darwin': + assert result.returncode == 2 and 'uses fp4 experts' in result.stderr and not record.exists(), result.stderr + assert bodies == [] and not (fp4 / 'mtp-q2_0.gguf').exists() + else: + assert result.returncode == 0 and bodies[0]['model'] == 'edge0_120b', result.stderr + config = json.loads((record / 'config.json').read_text()) + assert config['env']['STRATA_LAT_FUSED_Q3K'] == '0' and not (fp4 / 'scales').exists() + assert config['args'][1] == str(fp4) + assert_no_leftovers(record) print('PASS chat strata') diff --git a/edge0-infer/tests/integration/fake_strata.py b/edge0-infer/tests/integration/fake_strata.py index e4979867..dcd1cd73 100644 --- a/edge0-infer/tests/integration/fake_strata.py +++ b/edge0-infer/tests/integration/fake_strata.py @@ -41,6 +41,54 @@ ANSWER = 'Hello from fake Strata, answering in several pieces.' +# Stand-in for edge0-strata's tools/mtp_gguf.py. The delegate launcher runs it with the engine Python +# when a Metal template's prepare output is missing. FAKE_MTP_MODE selects a failure: +# fail exit 1 with an error on stderr and no output file +# exit0 exit 0 with no output file +# sleep record the call, then sleep (Ctrl-C must kill this process) +MTP_SCRIPT = '''"""Fake MTP conversion used by delegate-mode tests.""" +import json +import os +import sys +import time +from pathlib import Path + + +def main(): + args = sys.argv[1:] + out = None + if "--out" in args: + out = Path(args[args.index("--out") + 1]) + if out is not None: + log = out.parent / "prepare-calls.jsonl" + log.parent.mkdir(parents=True, exist_ok=True) + with log.open("a", encoding="utf-8") as handle: + handle.write(json.dumps({ + "argv": args, + "cwd": os.getcwd(), + "api_key": os.environ.get("STRATA_API_KEY"), + "pid": os.getpid(), + }) + "\\n") + mode = os.environ.get("FAKE_MTP_MODE", "") + if mode == "fail": + sys.stderr.write("mtp conversion failed\\n") + return 1 + if mode == "exit0": + return 0 + if mode == "sleep": + time.sleep(30) + return 0 + if out is None: + sys.stderr.write("missing --out\\n") + return 2 + out.write_bytes(b"mtp-q2_0") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) +''' + class Fake: def __init__(self, config, record): @@ -520,6 +568,31 @@ def call(method, path, body=None): if process.poll() is None: process.kill() process.wait(10) + with tempfile.TemporaryDirectory() as folder: + root = Path(folder) + script = root / 'mtp_gguf.py' + script.write_text(MTP_SCRIPT) + out = root / 'model' / 'mtp-q2_0.gguf' + + def run(mode): + env = dict(os.environ) + env.pop('STRATA_API_KEY', None) + if mode: + env['FAKE_MTP_MODE'] = mode + else: + env.pop('FAKE_MTP_MODE', None) + return subprocess.run([sys.executable, str(script), '--rt', str(root / 'mtp'), '--out', str(out)], + cwd=root, env=env, capture_output=True, text=True) + + ok = run('') + assert ok.returncode == 0 and out.read_bytes() == b'mtp-q2_0', (ok.returncode, ok.stderr) + recorded = json.loads((out.parent / 'prepare-calls.jsonl').read_text().splitlines()[0]) + assert Path(recorded['cwd']).resolve() == root.resolve() and recorded['pid'] > 0 and recorded['api_key'] is None + out.unlink() + failed = run('fail') + assert failed.returncode == 1 and 'mtp conversion failed' in failed.stderr and not out.exists() + silent = run('exit0') + assert silent.returncode == 0 and not out.exists() and silent.stderr == '' print('PASS fake strata: command line, ready file, status, streaming with timings, echo, records, failure, child cleanup') diff --git a/edge0-infer/tests/integration/serve_strata_test.py b/edge0-infer/tests/integration/serve_strata_test.py index 9fcaca5e..ff8fe525 100644 --- a/edge0-infer/tests/integration/serve_strata_test.py +++ b/edge0-infer/tests/integration/serve_strata_test.py @@ -194,7 +194,7 @@ def run(name, extra=(), engine_dir=engine, **fake): assert 'Python for edge0-strata not found' in text and not record.exists() other = make_engine(root / 'f-engine-2', contract='3') record, text = run('contract', engine_dir=other) - assert 'version 3 is not supported' in text and 'supports version 4' in text and not record.exists() + assert 'version 3 is not supported' in text and 'supports version 5' in text and not record.exists() record, text = run('start-failure', fail_start=5) assert 'exited before it was ready (exit code 5)' in text, text assert_no_leftovers(record) @@ -247,7 +247,11 @@ def proxy(cli, root, snapshots): assert initial['current'] is None and initial['last'] is None and initial['totals']['requests'] == 0 assert initial['instance']['pid'] == server.process.pid and initial['listen'] == {'host': '127.0.0.1', 'port': port, 'auth': False} assert initial['engine']['version'] == 'fake-strata 0.0' and initial['model']['context_length'] == 131072 - assert initial['engine']['devices'][0]['name'] == 'cuda:0' and initial['engine']['devices'][0]['description'] == 'Fake GPU' + device = initial['engine']['devices'][0] + expected_device = 'metal' if sys.platform == 'darwin' else 'cuda:0' + expected_backend = 'metal' if sys.platform == 'darwin' else 'cuda' + assert device['name'] == expected_device and device['description'] == 'Fake GPU' + assert device['backend'] == expected_backend and initial['model']['device'] == expected_device assert initial['memory'] == {'ledger': None, 'resident': None, 'physical': None} # Non-streaming chat: the request and the response bodies arrive byte for byte. @@ -680,20 +684,35 @@ def config_of(extra): return json.loads((record / 'config.json').read_text()) finally: assert server.stop() == 0 + mac = sys.platform == 'darwin' config = config_of([]) - assert config['gpu'] == 0 and config['model_name'] == ALIAS - assert config['env']['STRATA_LATENT_DENSE'] == str(model / 'latent_dense_q8.bin') and config['env']['STRATA_LATENT_DENSE_FMT'] == '8' + assert config['gpu'] == 0 and config['model_name'] == ALIAS and 'prepare' not in config + default_du = 'q4' if mac else 'q8' + assert config['env']['STRATA_LATENT_DENSE'] == str(model / f'latent_dense_{default_du}.bin') + assert config['env']['STRATA_LATENT_DENSE_FMT'] == default_du[1:] assert config['args'] == ['--pack', str(model), '--ple-io', 'direct', '--max-context', '131072'] assert Path(config['exe']) == engine / 'build' / ('strata.exe' if os.name == 'nt' else 'strata') and config['cwd'] == str(engine) - config = config_of(['--device', 'cuda:1', '--model-option', 'du=q4', '--model-option', 'ple-io=ram', '--context', '8192', - '--alias', 'big', '--prefill-device', 'cuda:1']) - assert config['gpu'] == 1 and config['model_name'] == 'big' - assert config['env']['STRATA_LATENT_DENSE'] == str(model / 'latent_dense_q4.bin') and config['env']['STRATA_LATENT_DENSE_FMT'] == '4' - assert config['args'] == ['--pack', str(model), '--ple-io', 'ram', '--max-context', '8192'] - assert config_of(['--device', 'cuda'])['gpu'] == 0 - for extra, message in ((['--device', 'cpu'], 'CUDA GPU'), (['--device', 'metal'], 'CUDA GPU'), (['--device', 'vulkan'], 'CUDA GPU'), - (['--device', 'cuda:x'], 'CUDA GPU'), (['--prefill-device', 'cpu'], 'must equal --device'), - (['--model-option', 'du=q5'], 'q8|q4'), (['--model-option', 'ple-io=mmap'], 'direct|ram'), + if mac: + config = config_of(['--device', 'metal', '--context', '8192', '--alias', 'big']) + assert config['gpu'] == 0 and config['model_name'] == 'big' and 'prepare' not in config + assert config['args'] == ['--pack', str(model), '--ple-io', 'direct', '--max-context', '8192'] + rejected = ((['--device', 'cpu'], 'must be metal'), (['--device', 'cuda'], 'must be metal'), + (['--device', 'cuda:0'], 'must be metal'), (['--device', 'cuda:1'], 'must be metal'), + (['--device', 'vulkan'], 'must be metal'), (['--prefill-device', 'cuda:0'], 'must be metal'), + (['--model-option', 'du=q8'], 'only supports du=q4'), (['--model-option', 'du=q5'], 'must be q4'), + (['--model-option', 'ple-io=ram'], 'only supports ple-io=direct'), + (['--model-option', 'ple-io=mmap'], 'must be direct')) + else: + config = config_of(['--device', 'cuda:1', '--model-option', 'du=q4', '--model-option', 'ple-io=ram', '--context', '8192', + '--alias', 'big', '--prefill-device', 'cuda:1']) + assert config['gpu'] == 1 and config['model_name'] == 'big' + assert config['env']['STRATA_LATENT_DENSE'] == str(model / 'latent_dense_q4.bin') and config['env']['STRATA_LATENT_DENSE_FMT'] == '4' + assert config['args'] == ['--pack', str(model), '--ple-io', 'ram', '--max-context', '8192'] + assert config_of(['--device', 'cuda'])['gpu'] == 0 + rejected = ((['--device', 'cpu'], 'CUDA GPU'), (['--device', 'metal'], 'CUDA GPU'), (['--device', 'vulkan'], 'CUDA GPU'), + (['--device', 'cuda:x'], 'CUDA GPU'), (['--prefill-device', 'cpu'], 'must equal --device'), + (['--model-option', 'du=q5'], 'q8|q4'), (['--model-option', 'ple-io=mmap'], 'direct|ram')) + for extra, message in (*rejected, (['--model-option', 'experts=resident'], 'unknown model option experts'), (['--prefill-chunk', '64'], '--prefill-chunk'), (['--temperature', '0'], '--temperature'), (['--max-tokens', '4'], '--max-tokens'), (['--no-warmup'], '--no-warmup'), @@ -752,19 +771,28 @@ def passthrough(cli, root): def model_formats(cli, root): - from strata_support import TEMPLATE + from strata_support import TEMPLATE, write_templates engine = make_engine(root / 'formats-engine') release = dict(TEMPLATE) release['args'] = [*TEMPLATE['args'], '--native', '${MODEL_DIR}/model-skeleton.gguf', '--native-dense-gguf', '${MODEL_DIR}/model-skeleton.gguf', '--expert-profile', '${MODEL_DIR}/profile.bin'] - for template in ('edge0-145b.json', 'edge0-120b.json'): - (engine / 'edge0' / 'models' / template).write_text(json.dumps(release)) - for identifier, du in (('edge0_145b', 'q8'), ('edge0_120b', 'q4')): + write_templates(engine, release) + mac = sys.platform == 'darwin' + for identifier, du in (('edge0_145b', 'q4' if mac else 'q8'), ('edge0_120b', 'q4')): for experts in ('q3k', 'fp4'): name = identifier.replace('_', '-') + '-' + experts model = make_model(root / name, model_name=name) record = root / (name + '-record') + if mac and experts == 'fp4': + result = subprocess.run([str(cli), 'serve', '--model', str(model), '--port', '0', *strata_args(engine)], + env=environment(record), capture_output=True, text=True, timeout=60) + assert result.returncode == 2 and 'uses fp4 experts' in result.stderr and 'Metal supports q3k' in result.stderr, result.stderr + assert not record.exists() and not (model / 'mtp-q2_0.gguf').exists() + inspected = subprocess.run([str(cli), 'inspect', '--model', str(model), '--json'], capture_output=True, text=True) + report = json.loads(inspected.stdout) + assert inspected.returncode == 1 and report['ok'] is False and 'uses fp4 experts' in inspected.stderr + continue server = Running(cli, model, engine, record) try: config = json.loads((record / 'config.json').read_text()) @@ -779,7 +807,7 @@ def model_formats(cli, root): assert config['env']['STRATA_LAT_FUSED_Q3K'] == ('1' if experts == 'q3k' else '0') assert Path(config['env']['STRATA_LATENT_DENSE']) == model / f'latent_dense_{du}.bin' assert config['env']['STRATA_LATENT_DENSE_FMT'] == du[1:] - assert 'log' not in config and config['exe'] != 'ignored' + assert 'log' not in config and 'prepare' not in config and config['exe'] != 'ignored' finally: assert server.stop() == 0 assert_no_leftovers(record) @@ -809,9 +837,12 @@ def model_formats(cli, root): def inspect_defaults(cli, root): import shutil + mac = sys.platform == 'darwin' keys = {'adapter', 'delegated', 'engine', 'directory', 'ok', 'files'} file_keys = {'filename', 'required', 'exists', 'bytes'} for identifier, default_du, other_du in (('edge0_145b', 'q8', 'q4'), ('edge0_120b', 'q4', 'q8')): + if mac: + default_du, other_du = 'q4', 'q8' model = make_model(root / ('inspect-' + identifier), model_name=identifier.replace('_', '-') + '-q3k') base = [str(cli), 'inspect', '--model', str(model), '--json'] @@ -827,12 +858,20 @@ def inspect(extra, ok): inspect([], True) (model / ('latent_dense_' + default_du + '.bin')).unlink() result, _ = inspect([], False) - assert result.stderr.rstrip().endswith('use --model-option du=' + other_du + ' instead.'), result.stderr - inspect(['--model-option', 'du=' + other_du], True) - result, _ = inspect(['--model-option', 'du=' + default_du], False) - suggested = result.stderr.rsplit(' use ', 1)[1].removesuffix(' instead.\n').split() - assert suggested == ['--model-option', 'du=' + other_du] - inspect(suggested, True) + if mac: + assert 'latent_dense_q4.bin' in result.stderr and '--model-option' not in result.stderr, result.stderr + rejected = subprocess.run([*base, '--model-option', 'du=' + other_du], capture_output=True, text=True) + report = json.loads(rejected.stdout) + assert rejected.returncode == 1 and report['ok'] is False and 'only supports du=q4' in rejected.stderr, rejected.stderr + (model / 'latent_dense_q4.bin').write_bytes(b'q4') + inspect([], True) + else: + assert result.stderr.rstrip().endswith('use --model-option du=' + other_du + ' instead.'), result.stderr + inspect(['--model-option', 'du=' + other_du], True) + result, _ = inspect(['--model-option', 'du=' + default_du], False) + suggested = result.stderr.rsplit(' use ', 1)[1].removesuffix(' instead.\n').split() + assert suggested == ['--model-option', 'du=' + other_du] + inspect(suggested, True) shutil.rmtree(model / 'scales') result, _ = inspect([], False) assert 'scales/' in result.stderr and '--model-option' not in result.stderr @@ -1005,12 +1044,126 @@ def launch_handshake(cli, root): print('PASS launch handshake: environment cleared, banner skipped, missing PID reported, no leftovers') +def mac_prepare(cli, root): + """Metal prepare and the Mac refusals. Other hosts skip these cases inside this process. + + Exiting 77 would skip the CUDA cases that share this ctest, and a separate test is outside this change's scope. + """ + if sys.platform != 'darwin': + print('SKIP macOS prepare: this host is not macOS') + return + engine = make_engine(root / 'mac-engine') + model = make_model(root / 'mac-model') + record = root / 'mac-first' + env = environment(record) + env['STRATA_API_KEY'] = 'secret-key' + server = Running(cli, model, engine, record, env=env) + try: + assert 'preparing' in server.text() and 'mtp-q2_0.gguf' in server.text() and '(one time)' in server.text() + assert (model / 'mtp-q2_0.gguf').read_bytes() == b'mtp-q2_0' + calls = [json.loads(line) for line in (model / 'prepare-calls.jsonl').read_text().splitlines()] + assert len(calls) == 1 and calls[0]['api_key'] is None, calls + assert Path(calls[0]['cwd']).resolve() == engine.resolve() + assert calls[0]['argv'][:2] == ['--rt', str(model / 'mtp')] + config = json.loads((record / 'config.json').read_text()) + assert 'prepare' not in config + status = json.loads(request(server.port, 'GET', '/status')[2]) + assert status['engine']['devices'][0]['name'] == 'metal' + assert status['engine']['devices'][0]['backend'] == 'metal' + assert status['model']['device'] == 'metal' + finally: + assert server.stop() == 0 + assert_no_leftovers(record) + rel_model = make_model(root / 'mac-rel-model') + rel_record = root / 'mac-rel' + rel_server = Running(cli, rel_model, engine, rel_record, extra=['--strata-python', os.path.relpath(sys.executable)]) + try: + assert (rel_model / 'mtp-q2_0.gguf').read_bytes() == b'mtp-q2_0' + rel_calls = [json.loads(line) for line in (rel_model / 'prepare-calls.jsonl').read_text().splitlines()] + assert Path(rel_calls[0]['cwd']).resolve() == engine.resolve() + finally: + assert rel_server.stop() == 0 + assert_no_leftovers(rel_record) + record = root / 'mac-second' + server = Running(cli, model, engine, record) + try: + assert 'preparing' not in server.text(), server.text() + assert len((model / 'prepare-calls.jsonl').read_text().splitlines()) == 1 + finally: + assert server.stop() == 0 + assert_no_leftovers(record) + + def rejected(name, extra, message, mode=None): + target = make_model(root / name, model_name='edge0-145b-q3k') + log = root / (name + '-record') + run_env = environment(log) + if mode: + run_env['FAKE_MTP_MODE'] = mode + result = subprocess.run([str(cli), 'serve', '--model', str(target), '--port', '0', *strata_args(engine), *extra], + capture_output=True, text=True, env=run_env, timeout=60) + assert result.returncode != 0 and 'Listening:' not in result.stderr, (name, result.returncode, result.stderr) + assert message in result.stderr, (name, result.stderr) + assert not log.exists(), ('a service process started', name) + return result + + failed = rejected('mac-fail', [], 'mtp conversion failed', mode='fail') + assert failed.returncode == 1 and 'mtp-q2_0.gguf' in failed.stderr and 'mtp_gguf.py' in failed.stderr, failed.stderr + assert 'output was not created' in failed.stderr + assert not (root / 'mac-fail' / 'edge0-145b' / 'mtp-q2_0.gguf').exists() + silent = rejected('mac-exit0', [], 'output was not created', mode='exit0') + assert silent.returncode == 1 and 'exit code 0' in silent.stderr and 'mtp_gguf.py' in silent.stderr, silent.stderr + assert not (root / 'mac-exit0' / 'edge0-145b' / 'mtp-q2_0.gguf').exists() + for name, extra, message in ( + ('mac-fp4', [], 'uses fp4 experts'), + ('mac-q8', ['--model-option', 'du=q8'], 'only supports du=q4'), + ('mac-ram', ['--model-option', 'ple-io=ram'], 'only supports ple-io=direct'), + ('mac-cuda', ['--device', 'cuda'], 'must be metal')): + if name == 'mac-fp4': + directory = make_model(root / name, model_name='edge0-120b-fp4') + log = root / (name + '-record') + result = subprocess.run([str(cli), 'serve', '--model', str(directory), '--port', '0', *strata_args(engine)], + capture_output=True, text=True, env=environment(log), timeout=60) + assert result.returncode == 2 and message in result.stderr and not log.exists(), result.stderr + assert not (directory / 'mtp-q2_0.gguf').exists() + else: + result = rejected(name, extra, message) + assert result.returncode == 2, (name, result.returncode, result.stderr) + assert not (root / name / 'edge0-145b' / 'mtp-q2_0.gguf').exists() + # Ctrl-C while the script is still running kills it and does not start the service. + sleeping = make_model(root / 'mac-sleep') + log = root / 'mac-sleep-record' + run_env = environment(log) + run_env['FAKE_MTP_MODE'] = 'sleep' + process = subprocess.Popen([str(cli), 'serve', '--model', str(sleeping), '--port', '0', *strata_args(engine)], + stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True, env=run_env) + calls = sleeping / 'prepare-calls.jsonl' + deadline = time.time() + 15 + while time.time() < deadline and not calls.exists(): + if process.poll() is not None: + raise AssertionError(process.stderr.read()) + time.sleep(0.05) + assert calls.exists(), 'prepare script did not start' + pid = json.loads(calls.read_text().splitlines()[0])['pid'] + process.send_signal(signal.SIGINT) + try: + _, stderr = process.communicate(timeout=15) + except subprocess.TimeoutExpired: + process.kill() + process.communicate(timeout=10) + raise AssertionError('Ctrl-C during prepare did not stop edge0-infer') + assert process.returncode == 1 and 'interrupted while preparing' in stderr, (process.returncode, stderr) + assert not log.exists() and not (sleeping / 'mtp-q2_0.gguf').exists() + assert gone(pid), 'prepare script outlived Ctrl-C' + print('PASS macOS prepare: once, then skipped; failure and exit 0 leave no file or service; fp4, du=q8, ple-io=ram and cuda rejected; Ctrl-C kills the script') + + def main(): parser = argparse.ArgumentParser() parser.add_argument('--cli', required=True) args = parser.parse_args() with tempfile.TemporaryDirectory() as folder: root = Path(folder) + mac_prepare(args.cli, root) launch_handshake(args.cli, root) interrupted_responses(args.cli, root) inspect_defaults(args.cli, root) diff --git a/edge0-infer/tests/integration/strata_support.py b/edge0-infer/tests/integration/strata_support.py index 48659919..ea7b1ae3 100644 --- a/edge0-infer/tests/integration/strata_support.py +++ b/edge0-infer/tests/integration/strata_support.py @@ -27,16 +27,37 @@ } # The desktop App's readiness rule (examples/desktop/src-tauri/src/service.rs, ready_url). APP_READY = re.compile(r'^Listening: http://(?:127\.0\.0\.1|0\.0\.0\.0):(\d+) model=(\S.*)$') - - -def make_engine(root, contract='4'): +# Only the Metal template declares this. The fake tools/mtp_gguf.py writes the output unless FAKE_MTP_MODE says otherwise. +PREPARE = [{ + 'output': '${MODEL_DIR}/mtp-q2_0.gguf', + 'script': 'tools/mtp_gguf.py', + 'args': ['--rt', '${MODEL_DIR}/mtp', '--out', '${MODEL_DIR}/mtp-q2_0.gguf'], +}] + + +def write_templates(engine, body): + """CUDA templates as given; Metal templates are the same body plus the prepare step.""" + models = Path(engine) / 'edge0' / 'models' + models.mkdir(parents=True, exist_ok=True) + metal = dict(body) + metal['prepare'] = PREPARE + text = json.dumps(body, indent=1) + metal_text = json.dumps(metal, indent=1) + for name in ('edge0-145b.json', 'edge0-120b.json'): + (models / name).write_text(text) + (models / name.replace('.json', '-metal.json')).write_text(metal_text) + + +def make_engine(root, contract='5'): + from fake_strata import MTP_SCRIPT engine = Path(root) / 'edge0-strata' (engine / 'serve').mkdir(parents=True) - (engine / 'edge0' / 'models').mkdir(parents=True) + (engine / 'tools').mkdir(parents=True) shutil.copy(HERE / 'fake_strata.py', engine / 'serve' / 'server.py') + (engine / 'tools' / 'mtp_gguf.py').write_text(MTP_SCRIPT) + (engine / 'edge0').mkdir(parents=True) (engine / 'edge0' / 'CONTRACT').write_text(contract + '\n') - (engine / 'edge0' / 'models' / 'edge0-145b.json').write_text(json.dumps(TEMPLATE, indent=1)) - (engine / 'edge0' / 'models' / 'edge0-120b.json').write_text(json.dumps(TEMPLATE, indent=1)) + write_templates(engine, TEMPLATE) return engine diff --git a/edge0-infer/tests/platform/strata/check_mtp_gguf.py b/edge0-infer/tests/platform/strata/check_mtp_gguf.py new file mode 100644 index 00000000..1cb54931 --- /dev/null +++ b/edge0-infer/tests/platform/strata/check_mtp_gguf.py @@ -0,0 +1,114 @@ +"""Compare a rebuilt MTP GGUF with the original mtp-q2_0.gguf. + +Same-type tensors must match byte for byte. A Q8_0 tensor rebuilt from the runtime files is compared with the +original BF16 tensor quantized the way the Metal engine does at load (ref::quantize_q8_0: fp32 d = amax / 127, +fp16 d by round-to-nearest-even, roundf of x * (1 / d)). +""" +import struct +import sys + +import numpy as np + + +def read(path): + f = open(path, "rb") + magic, ver, nt, nkv = struct.unpack("<4sIQQ", f.read(24)) + assert magic == b"GGUF" and ver == 3 + + def rs(): + n, = struct.unpack(" Q8_0 vs engine quantization: {bad_blocks}/{ref.shape[0]} blocks differ " + f"(scales {bad_scales}, values {int((qd > 0).sum())}, max step {int(qd.max())})") + if bad_blocks > 0: + ok = False + else: + print(f"{name}: type {ot} vs {nt}") + ok = False + print("RESULT:", "same-type tensors identical" if ok else "MISMATCH") + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/edge0-infer/tests/platform/strata/compare_outputs.py b/edge0-infer/tests/platform/strata/compare_outputs.py index a2468489..30d539d3 100644 --- a/edge0-infer/tests/platform/strata/compare_outputs.py +++ b/edge0-infer/tests/platform/strata/compare_outputs.py @@ -54,8 +54,11 @@ def main(): parser.add_argument('--device', default='cuda:0') parser.add_argument('--strata-dir') parser.add_argument('--max-tokens', type=int, default=200) + parser.add_argument('--question', action='append', default=[], + help='ask these instead of the built-in three; repeat the option for each question') parser.add_argument('--out', help='write all answers as JSON') args = parser.parse_args() + questions = args.question or QUESTIONS command = [args.cli, 'serve', '--model', args.model, '--device', args.device, '--port', '0'] if args.strata_dir: command += ['--strata-dir', args.strata_dir] @@ -81,7 +84,7 @@ def read(): if 'proxy' not in ports or 'direct' not in ports: raise SystemExit('edge0-infer serve did not become ready') results = [] - for question in QUESTIONS: + for question in questions: first = ask(ports['direct'], ports['alias'], question, args.max_tokens) second = ask(ports['direct'], ports['alias'], question, args.max_tokens) proxied = ask(ports['proxy'], ports['alias'], question, args.max_tokens) diff --git a/edge0-infer/tests/unit/cli/strata_registry_test.cpp b/edge0-infer/tests/unit/cli/strata_registry_test.cpp index 93e7e324..96e7e233 100644 --- a/edge0-infer/tests/unit/cli/strata_registry_test.cpp +++ b/edge0-infer/tests/unit/cli/strata_registry_test.cpp @@ -114,7 +114,8 @@ int main() { } // Devices and model options. - CHECK(parse_device("cuda") == 0 && parse_device("cuda:0") == 0 && parse_device("cuda:1") == 1 && parse_device("cuda:12") == 12); + CHECK(parse_device("cuda").gpu == 0 && parse_device("cuda").name == "cuda:0" && parse_device("cuda:0").gpu == 0 && + parse_device("cuda:1").gpu == 1 && parse_device("cuda:1").name == "cuda:1" && parse_device("cuda:12").gpu == 12); for (const auto* bad : {"cpu", "metal", "vulkan", "cuda:", "cuda:x", "cuda:-1", "CUDA", "cuda:1 "}) CHECK(throws_config([&] { parse_device(bad); }, "CUDA GPU")); const auto defaults = resolve_options(registration, {}); @@ -206,6 +207,45 @@ int main() { CHECK(given.strata_args.size() == 2); } + // prepare: placeholders expand like every other field, a relative script is under the engine directory, + // and the field is not part of the config handed to server.py. + { + const auto with_prepare = Json::parse(R"({"exe":"${STRATA_DIR}/build/strata","args":["--mtp","${MODEL_DIR}/mtp"], + "prepare":[{"output":"${MODEL_DIR}/mtp-q2_0.gguf","script":"tools/mtp_gguf.py", + "args":["--rt","${MODEL_DIR}/mtp","--out","${MODEL_DIR}/mtp-q2_0.gguf"]}]})"); + std::vector steps; + const auto expanded = expand_template(with_prepare, registration, options, Platform::macos, &steps); + CHECK(!expanded.contains("prepare")); + // The template keeps the slash after ${MODEL_DIR}; the expanded text is not re-normalized. + const auto mtp = options.model.string() + "/mtp"; + const auto mtp_file = options.model.string() + "/mtp-q2_0.gguf"; + CHECK(expanded["args"][1] == mtp); + CHECK(steps.size() == 1); + CHECK(steps[0].output == fs::path(mtp_file)); + CHECK(steps[0].script == options.strata_dir / "tools" / "mtp_gguf.py"); + CHECK((steps[0].args == std::vector{"--rt", mtp, "--out", mtp_file})); + auto absolute = with_prepare; + absolute["prepare"][0]["script"] = "${STRATA_DIR}/tools/mtp_gguf.py"; + expand_template(absolute, registration, options, Platform::macos, &steps); + CHECK(steps.size() == 1 && steps[0].script == options.strata_dir / "tools" / "mtp_gguf.py"); + expand_template(Json::parse(R"({"exe":"x","args":[]})"), registration, options, Platform::posix, &steps); + CHECK(steps.empty()); + const auto stripped = expand_template(with_prepare, registration, options, Platform::posix); + CHECK(!stripped.contains("prepare")); + CHECK(throws_config([&] { + expand_template(Json::parse(R"({"exe":"x","args":[],"prepare":[{"output":"${NOPE}","script":"t.py","args":[]}]})"), + registration, options, Platform::posix, &steps); + }, "${NOPE}")); + CHECK(throws_config([&] { + expand_template(Json::parse(R"({"exe":"x","args":[],"prepare":{"output":"a"}})"), + registration, options, Platform::posix, &steps); + }, "prepare must be an array")); + CHECK(throws_config([&] { + expand_template(Json::parse(R"({"exe":"x","args":[],"prepare":[{"output":"","script":"t.py","args":[]}]})"), + registration, options, Platform::posix, &steps); + }, "needs output, script and args")); + } + // The release templates: every model/format resolves all model files and D/U choices. CHECK(registry()[1].id == "edge0_120b"); CHECK(registry()[1].config_template == "edge0/models/edge0-120b.json"); @@ -288,11 +328,13 @@ int main() { CHECK(throws_config([&] { locate_engine(engine.string(), ""); }, "not a directory")); fs::create_directories(engine / "edge0"); CHECK(throws_config([&] { locate_engine(engine.string(), ""); }, "edge0/CONTRACT")); - write(engine / "edge0" / "CONTRACT", "4\n"); + write(engine / "edge0" / "CONTRACT", "5\n"); CHECK(locate_engine(engine.string(), "") == fs::absolute(engine)); check_contract(engine); write(engine / "edge0" / "CONTRACT", "3\n"); - CHECK(throws_config([&] { check_contract(engine); }, "version 3 is not supported; this edge0-infer supports version 4")); + CHECK(throws_config([&] { check_contract(engine); }, "version 3 is not supported; this edge0-infer supports version 5")); + write(engine / "edge0" / "CONTRACT", "4\n"); + CHECK(throws_config([&] { check_contract(engine); }, "version 4 is not supported; this edge0-infer supports version 5")); write(engine / "edge0" / "CONTRACT", "one\n"); CHECK(throws_config([&] { check_contract(engine); }, "unreadable")); @@ -321,7 +363,7 @@ int main() { platform_args.strata_dir = engine.string(); const auto cuda = t.path / "CUDA toolkit"; const auto platform_template = Json::parse(R"({"exe":"${STRATA_DIR}/build/strata${EXE_SUFFIX}","args":[],"lib_dirs":["${CUDA_LIB_DIR}"]})"); - for (const auto platform : {Platform::posix, Platform::windows}) { + for (const auto platform : {Platform::posix, Platform::windows, Platform::macos}) { fs::remove_all(cuda); for (int layout = 0; layout < 4; ++layout) { if (layout == 2) fs::create_directories(cuda / "bin"); @@ -329,6 +371,10 @@ int main() { const auto launch = resolve(platform_args, *detect(model), installed_exe, platform, layout == 0 ? "" : cuda.string()); const bool windows = platform == Platform::windows; + CHECK(launch.device == default_device(platform)); + CHECK(launch.gpu == 0); + CHECK(launch.model_options.at("du") == (platform == Platform::macos ? "q4" : "q8")); + CHECK(launch.model_options.at("ple-io") == "direct"); CHECK(launch.python == engine / ".venv" / (windows ? fs::path("Scripts") / "python.exe" : fs::path("bin") / "python")); const auto expanded = expand_template(platform_template, registration, launch, platform); CHECK(fs::path(expanded["exe"].get()) == engine / "build" / (windows ? "strata.exe" : "strata")); @@ -354,20 +400,110 @@ int main() { CHECK(throws_config([&] { expand_template(Json::parse(bad), registration, launch, platform); }, "${CUDA_LIB_DIR}")); } platform_args.python = (t.path / "custom python").string(); - CHECK(resolve(platform_args, *detect(model), installed_exe, platform, "").python == *platform_args.python); + CHECK(resolve(platform_args, *detect(model), installed_exe, platform, "").python == fs::absolute(*platform_args.python)); + platform_args.python = "rel-python"; + CHECK(resolve(platform_args, *detect(model), installed_exe, platform, "").python == fs::absolute("rel-python")); platform_args.python.reset(); CHECK(resolve(platform_args, *detect(model), installed_exe, platform, "").strata_args.empty()); platform_args.strata_args = {"--pool-workers", "24"}; CHECK((resolve(platform_args, *detect(model), installed_exe, platform, "").strata_args == platform_args.strata_args)); platform_args.strata_args.clear(); - for (const auto* version : {"0", "1", "2", "3"}) { + for (const auto* version : {"0", "1", "2", "3", "4"}) { write(engine / "edge0" / "CONTRACT", version); - CHECK(throws_config([&] { check_contract(engine); }, "supports version 4")); + CHECK(throws_config([&] { check_contract(engine); }, "supports version 5")); } - write(engine / "edge0" / "CONTRACT", "4\n"); + write(engine / "edge0" / "CONTRACT", "5\n"); check_contract(engine); } + // Device rules for every platform parameter: defaults, rejection, the Metal template, and inspect. + { + auto run_inspect = [&](Platform platform, const std::map& given) { + return inspect_report(*detect(model), model, given, default_device(platform)); + }; + const auto metal_template = fs::path("edge0") / "models" / "edge0-145b-metal.json"; + write(engine / registration.config_template, R"({"exe":"cuda-template","args":[]})"); + write(engine / metal_template, R"({"exe":"metal-template","args":["--max-context","32768"]})"); + for (const auto platform : {Platform::posix, Platform::windows, Platform::macos}) { + CHECK(default_device(platform) == (platform == Platform::macos ? "metal" : "cuda:0")); + if (platform == Platform::macos) { + CHECK(parse_device("metal", platform).name == "metal" && parse_device("metal", platform).gpu == 0); + for (const auto* bad : {"cuda", "cuda:0", "cuda:1"}) + CHECK(throws_config([&] { parse_device(bad, platform); }, "must be metal")); + } else { + CHECK(throws_config([&] { parse_device("metal", platform); }, "CUDA GPU")); + } + complete(model); + write(engine / "edge0" / "CONTRACT", "5\n"); + CommonArgs bare; + bare.model = model.string(); + bare.strata_dir = engine.string(); + const auto launch = resolve(bare, *detect(model), installed_exe, platform, ""); + CHECK(launch.device == default_device(platform)); + CHECK(load_template(engine, registration, launch.device)["exe"] == (platform == Platform::macos ? "metal-template" : "cuda-template")); + if (platform == Platform::macos) { + CHECK(launch.model_options.at("du") == "q4" && launch.model_options.at("ple-io") == "direct"); + write(model / "model.json", R"({"model_name":"edge0-120b-q3k"})"); + CHECK(resolve(bare, *detect(model), installed_exe, platform, "").model_options.at("du") == "q4"); + write(model / "model.json", R"({"model_name":"edge0-145b-q3k"})"); + fs::remove(model / "experts.bin"); + fs::remove(model / "latent_dense_q4.bin"); + write(model / "model.json", R"({"model_name":"edge0-145b-fp4"})"); + fs::remove_all(model / "scales"); + CHECK(throws_config([&] { resolve(bare, *detect(model), installed_exe, platform, ""); }, + "unsupported on Metal: edge0-145b-fp4 uses fp4 experts; Metal supports q3k")); + CommonArgs illegal = bare; + illegal.model_options = {{"du", "q5"}}; + CHECK(throws_config([&] { resolve(illegal, *detect(model), installed_exe, platform, ""); }, "must be")); + CHECK(throws_config([&] { run_inspect(platform, {{"du", "q5"}}); }, "must be")); + auto result = run_inspect(platform, {}); + CHECK(result.code == 1 && result.message.find("uses fp4 experts") != std::string::npos && + result.message.find("Metal supports q3k") != std::string::npos); + CHECK(result.report["ok"] == false && result.report["files"].is_array() && !result.report["files"].empty()); + complete(model); + CommonArgs rejected = bare; + rejected.model_options = {{"du", "q8"}}; + fs::remove(model / "latent_dense_q4.bin"); + CHECK(throws_config([&] { resolve(rejected, *detect(model), installed_exe, platform, ""); }, + "unsupported on Metal: du=q8; Metal only supports du=q4")); + result = run_inspect(platform, {{"du", "q8"}}); + CHECK(result.code == 1 && result.report["ok"] == false && result.message.find("only supports du=q4") != std::string::npos); + rejected.model_options = {{"ple-io", "ram"}}; + CHECK(throws_config([&] { resolve(rejected, *detect(model), installed_exe, platform, ""); }, + "Metal only supports ple-io=direct")); + result = run_inspect(platform, {{"ple-io", "ram"}}); + CHECK(result.code == 1 && result.report["ok"] == false && result.message.find("ple-io=direct") != std::string::npos); + CHECK(throws_config([&] { run_inspect(platform, {{"experts", "q3k"}}); }, "unknown model option experts")); + write(model / "latent_dense_q4.bin", "latent_dense_q4.bin"); + result = run_inspect(platform, {}); + CHECK(result.code == 0 && result.report["ok"] == true && result.message.empty()); + fs::remove(engine / metal_template); + CHECK(throws_config([&] { load_template(engine, registration, "metal"); }, "edge0-145b-metal.json")); + write(engine / metal_template, R"({"exe":"metal-template","args":[]})"); + } else { + CommonArgs prefill = bare; + prefill.prefill_device = "cpu"; + CHECK(throws_config([&] { resolve(prefill, *detect(model), installed_exe, platform, ""); }, "must equal")); + prefill.device = "cuda"; + prefill.prefill_device = "cuda:0"; + CHECK(throws_config([&] { resolve(prefill, *detect(model), installed_exe, platform, ""); }, "must equal")); + prefill.prefill_device = "cuda:01"; + CHECK(throws_config([&] { resolve(prefill, *detect(model), installed_exe, platform, ""); }, "must equal")); + CommonArgs cuda = bare; + cuda.device = "cuda:1"; + CHECK(resolve(cuda, *detect(model), installed_exe, platform, "").device == "cuda:1"); + cuda.model_options = {{"du", "q8"}}; + CHECK(resolve(cuda, *detect(model), installed_exe, platform, "").model_options.at("du") == "q8"); + write(model / "model.json", R"({"model_name":"edge0-120b-fp4"})"); + fs::remove_all(model / "scales"); + CHECK(resolve(bare, *detect(model), installed_exe, platform, "").experts == "fp4"); + const auto result = run_inspect(platform, {}); + CHECK(result.code == 0 && result.report["ok"] == true); + } + } + complete(model); + } + // Response accounting: plain JSON, SSE split at arbitrary points, other dialects. { bool active = true;int finished = 0; @@ -622,7 +758,14 @@ int main() { input.active = ActiveRequest{"/v1/messages", "127.0.0.1", input.now - std::chrono::seconds(2), Accounting{}}; auto status = status_json(input); CHECK(status["engine"]["version"] == "0.1.55" && status["engine"]["ggml_version"].is_null() && status["engine"]["backends"].is_null()); - CHECK(status["engine"]["devices"][0]["name"] == "cuda:1" && status["engine"]["devices"][0]["memory_total_bytes"] == 4096LL * 1048576); + CHECK(status["engine"]["devices"][0]["name"] == "cuda:1" && status["engine"]["devices"][0]["backend"] == "cuda" && + status["engine"]["devices"][0]["memory_total_bytes"] == 4096LL * 1048576); + input.device = "metal"; + status = status_json(input); + CHECK(status["engine"]["devices"][0]["name"] == "metal" && status["engine"]["devices"][0]["backend"] == "metal" && + status["model"]["device"] == "metal"); + input.device = "cuda:1"; + status = status_json(input); CHECK(status["engine"]["devices"][0]["memory_free_bytes"] == 3072LL * 1048576 && status["engine"]["devices"][0]["description"] == "RTX"); CHECK(status["model"]["context_length"] == 131072 && status["model"]["device"] == "cuda:1" && status["model"]["prefill_device"] == "cuda:1"); CHECK(status["memory"]["ledger"].is_null() && status["memory"]["resident"].is_null() && status["memory"]["physical"].is_null()); diff --git a/edge0-infer/tools/checks/check_layers.py b/edge0-infer/tools/checks/check_layers.py index 41d1a18b..dfcbe130 100644 --- a/edge0-infer/tools/checks/check_layers.py +++ b/edge0-infer/tools/checks/check_layers.py @@ -10,10 +10,12 @@ _spec.loader.exec_module(_common) cmake_commands, files = _common.cmake_commands, _common.files -# Windows console/stdio setup and delegated process management outside src/platform/. -# Only these files, and only these headers. +# Platform headers used by delegated process management, outside src/platform/. +# tools/cli/serve_delegate_process.cpp: Windows process APIs, and on macOS the +# header that locates this executable so a sibling edge0-strata/ can be found +# when the program is started through PATH. Only these files, and only these headers. CONSOLE_SETUP = { - "tools/cli/serve_delegate_process.cpp": {"windows.h"}, + "tools/cli/serve_delegate_process.cpp": {"windows.h", "mach-o/dyld.h", "spawn.h", "unistd.h", "fcntl.h", "sys/wait.h"}, "tools/cli/main.cpp": {"windows.h", "io.h", "fcntl.h"}, "tests/unit/app/template_tools_probe.cpp": {"io.h", "fcntl.h"}, } diff --git a/edge0-infer/tools/cli/serve_delegate.h b/edge0-infer/tools/cli/serve_delegate.h index 805e752a..2deefe10 100644 --- a/edge0-infer/tools/cli/serve_delegate.h +++ b/edge0-infer/tools/cli/serve_delegate.h @@ -22,7 +22,7 @@ namespace delegate { using Json = nlohmann::ordered_json; // The launch contract version this edge0-infer understands (edge0-strata/edge0/CONTRACT). -constexpr int supported_contract = 4; +constexpr int supported_contract = 5; struct RegisteredFile { std::string path; // relative to the model directory @@ -37,6 +37,8 @@ struct Registration { std::vector>> options; // name, allowed values (first = default) bool thinking = false, thinking_default = false; std::map names; // model_name -> expert format + // Allowed values for one device name. A device that is absent uses `options`. The first value is the default. + std::map>>> device_options; }; const std::vector& registry(); @@ -53,29 +55,47 @@ struct FileState { struct Match { const Registration* registration = nullptr; std::vector files; + std::string model_name; // the model.json value that selected this registration std::string experts; // selected by model_name, never a user option // Required files (for the given options) that are absent. std::vector missing(const std::map& options) const; // Inspect requires the selected configuration, using the same defaults as serve for omitted options. - bool complete(const std::map& selected = {}) const; + // `device` is "metal" or "cuda:"; the default is the CUDA device, so a caller that has not chosen a platform + // keeps the CUDA rules. inspect passes the host's default device. + bool complete(const std::map& selected = {}, const std::string& device = "cuda:0") const; // Option choices ("