diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 5f64f33..e4bffa9 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -61,6 +61,8 @@ jobs: run: | cp -r build/_deps/llama-src /tmp/llama-patchtest python3 scripts/patch_ggml_cuda_ext_hook.py /tmp/llama-patchtest + python3 scripts/patch_ggml_sycl_ext_hook.py /tmp/llama-patchtest + python3 scripts/patch_ggml_sycl_ext_hook.py /tmp/llama-patchtest # idempotent python3 scripts/patch_ggml_openvino.py /tmp/llama-patchtest python3 scripts/patch_ggml_openvino.py /tmp/llama-patchtest # idempotent @@ -88,6 +90,89 @@ jobs: -DCMAKE_EXE_LINKER_FLAGS=-Wl,--allow-shlib-undefined -DVLA_BUILD_TESTS=ON -DVLA_WERROR=ON cmake --build build -j"$(nproc)" + # Intel GPUs through oneAPI SYCL. The runner has no GPU, so ggml-sycl itself + # cannot start; ctest still runs everything, and test_foldquant_sycl_op drives + # the FoldQuant kernels straight through the extension hook on the OpenCL CPU + # device that the DPC++ compiler package ships, against the CPU reference. + build-sycl: + runs-on: ubuntu-24.04 + steps: + - name: free disk (oneAPI + a SYCL build need ~15 GB) + run: sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL + - uses: actions/checkout@v7 + - name: deps + run: | + wget -qO- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB \ + | sudo gpg --dearmor -o /usr/share/keyrings/oneapi-archive-keyring.gpg + echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" \ + | sudo tee /etc/apt/sources.list.d/oneAPI.list + sudo apt-get update -qq + sudo apt-get install -y -qq --no-install-recommends \ + build-essential cmake git ca-certificates pkg-config python3 \ + libzmq3-dev cppzmq-dev libprotobuf-dev protobuf-compiler \ + intel-oneapi-compiler-dpcpp-cpp-2026.1 intel-oneapi-mkl-devel-2026.1 intel-oneapi-dnnl-devel-2026.0 + - name: read llama.cpp pin + id: pin + run: echo "tag=$(bash scripts/llama_tag.sh)" >> "$GITHUB_OUTPUT" + - uses: actions/cache@v6 + with: + path: build/_deps + key: llama-${{ steps.pin.outputs.tag }}-${{ runner.os }}-sycl-${{ hashFiles('scripts/patch_ggml_sycl_ext_hook.py') }} + - name: build + ctest (OpenCL CPU device) + shell: bash + env: + ONEAPI_DEVICE_SELECTOR: opencl:cpu + VLA_FQ_TEST_REQUIRE_DEVICE: 1 + run: | + source /opt/intel/oneapi/setvars.sh + sycl-ls + cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_SYCL=ON \ + -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DVLA_BUILD_TESTS=ON -DVLA_WERROR=ON + cmake --build build -j"$(nproc)" + ctest --test-dir build --output-on-failure + + # Intel CPUs, GPUs and NPUs through OpenVINO. The CPU plugin runs on the + # runner, so test_foldquant_ov_op executes the translated FoldQuant graph and + # checks it against the CPU reference. + build-openvino: + runs-on: ubuntu-24.04 + env: + OV_ARCHIVE: openvino_toolkit_ubuntu24_2026.4.0.22959.99c81491cc3_x86_64 + steps: + - uses: actions/checkout@v7 + - name: deps + run: | + sudo apt-get update -qq + sudo apt-get install -y -qq --no-install-recommends \ + build-essential cmake git ca-certificates pkg-config python3 \ + libzmq3-dev cppzmq-dev libprotobuf-dev protobuf-compiler \ + opencl-clhpp-headers ocl-icd-opencl-dev opencl-headers + - uses: actions/cache@v6 + id: ov + with: + path: /opt/openvino + key: ${{ env.OV_ARCHIVE }} + - name: OpenVINO runtime + if: steps.ov.outputs.cache-hit != 'true' + run: | + curl -fsSL "https://storage.openvinotoolkit.org/repositories/openvino/packages/2026.4/linux/${OV_ARCHIVE}.tgz" \ + | tar xz -C /tmp + sudo mv "/tmp/${OV_ARCHIVE}" /opt/openvino + - name: read llama.cpp pin + id: pin + run: echo "tag=$(bash scripts/llama_tag.sh)" >> "$GITHUB_OUTPUT" + - uses: actions/cache@v6 + with: + path: build/_deps + key: llama-${{ steps.pin.outputs.tag }}-${{ runner.os }}-openvino-${{ hashFiles('scripts/patch_ggml_openvino.py') }} + - name: build + ctest (CPU plugin) + shell: bash + run: | + source /opt/openvino/setupvars.sh + cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_OPENVINO=ON -DVLA_BUILD_TESTS=ON -DVLA_WERROR=ON + cmake --build build -j"$(nproc)" + ctest --test-dir build --output-on-failure + build-backends: runs-on: ${{ matrix.os }} strategy: diff --git a/CHANGELOG.md b/CHANGELOG.md index 5164360..07b6029 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,30 @@ Notable changes to vla.cpp. Format loosely follows [Keep a Changelog](https://ke ## [Unreleased] +### Added + +- **FoldQuant W8A8 / W4A4 inference** for GR00T N1.5 / N1.6 / N1.7 and π0.5. A + FoldQuant GGUF carries the language backbone and the action + module as INT8 or INT4 codes with per-row scales in a block-Hadamard, + SmoothQuant-folded frame; activations are quantized per token and the + projections run as integer GEMMs (`src/kernels/foldquant/` on CUDA, + `src/sycl/vla_sycl_foldquant.cpp` on SYCL with oneDNN's int8 matmul, both + bit-identical to the reference in `src/foldquant_ref.cpp` that the CPU runs; + OpenVINO ops on the OpenVINO CPU and GPU plugins with INT8/INT4 weight + constants). On every other backend + (Metal, Hexagon, OpenCL, the OpenVINO NPU) the sites are read back as + float weights with FoldQuant's rounding and run the stock float path; + `VLA_FQ_DEQUANT=1` does the same on CUDA or CPU. The format and arithmetic are + in `docs/QUANTIZATION.md`. +- `scripts/convert_quantized_model_to_gguf.py` converts a FoldQuantVLA quantized + model (its quantized checkpoint or its earlier fake-quant state) to a FoldQuant + GGUF with no calibration and nothing re-rounded; `--check-onnx` byte-compares + every site against the TensorRT plugin graphs. The family converters expose + `convert(ckpt, out, writer_factory=...)` for it. +- `scripts/foldquant_fake_export.py` (uncalibrated file for bring-up), + `scripts/inspect_gguf_quant.py` (contract check) and `scripts/foldquant_ref.py` + (numpy reference). + ## [0.4.0] - 2026-09-30 ### Added diff --git a/CMakeLists.txt b/CMakeLists.txt index ce078c1..2a01c88 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -51,6 +51,12 @@ set(LLAMA_BUILD_SERVER OFF CACHE BOOL "" FORCE) set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE) set(LLAMA_BUILD_TESTS OFF CACHE BOOL "" FORCE) set(_vla_llama_patch "") +if(GGML_SYCL AND NOT GGML_CUDA) + # The hook FoldQuant's SYCL kernels (src/sycl/vla_sycl_foldquant.cpp) run behind. + find_package(Python3 COMPONENTS Interpreter REQUIRED) + set(_vla_llama_patch PATCH_COMMAND ${Python3_EXECUTABLE} + ${CMAKE_CURRENT_SOURCE_DIR}/scripts/patch_ggml_sycl_ext_hook.py ) +endif() if(GGML_CUDA) find_package(Python3 COMPONENTS Interpreter REQUIRED) set(_vla_llama_patch PATCH_COMMAND ${Python3_EXECUTABLE} @@ -117,6 +123,17 @@ FetchContent_Declare(llama ) FetchContent_MakeAvailable(llama) +# FoldQuant on OpenVINO: the translator for FoldQuant's two custom nodes is +# in-tree and compiled into ggml's OpenVINO backend, which reaches it through the +# hook scripts/patch_ggml_openvino.py adds (op table entry, supports_op). +if(GGML_OPENVINO) + target_sources(ggml-openvino PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src/openvino/foldquant_ov.cpp) + target_include_directories(ggml-openvino PRIVATE + ${llama_SOURCE_DIR}/ggml/src/ggml-openvino + ${llama_SOURCE_DIR}/ggml/src + ${CMAKE_CURRENT_SOURCE_DIR}/src) +endif() + # Drop the whole fetched tree from the default build: llama.cpp's ~20 tool # binaries and their impl libraries are dead weight here. What we link (llama, # ggml, mtmd, llama-common) is still built, pulled in as a dependency. Walking @@ -189,6 +206,15 @@ if(VLA_SPM) endif() endif() +# oneAPI's icpx (the compiler of a SYCL build) defaults to -fp-model=fast, which +# reassociates sums and drops IEEE division and NaN semantics; vla.cpp's code and +# tests are written for what GCC, Clang and MSVC do by default. Set here, after +# the dependencies, so ggml keeps its own flags. The FoldQuant sources then add +# -ffp-contract=off on top, which icpx reports as an override; that is the intent. +if(CMAKE_CXX_COMPILER_ID STREQUAL "IntelLLVM") + add_compile_options(-fp-model=precise -Wno-overriding-option) +endif() + # llama.cpp turns BUILD_SHARED_LIBS on, so these follow it. That is harmless on # ELF, which exports everything, but a Windows DLL exports nothing unannotated: # the import library comes out empty and every consumer fails to link. Nothing @@ -203,6 +229,8 @@ add_library(vla_core ${VLA_CORE_LIB_TYPE} src/model.cpp src/loader.cpp src/options.cpp + src/foldquant.cpp + src/foldquant_ref.cpp src/modules/action_expert.cpp src/modules/dit_head.cpp src/modules/encoder.cpp @@ -231,6 +259,10 @@ target_include_directories(vla_core ) # The VLA archs call no llama_* API; only vlm_core needs llama. target_link_libraries(vla_core PUBLIC ggml) +# The FoldQuant CPU reference must match the CUDA and SYCL kernels bit for bit, +# so no FMA contraction on either side (aarch64 GCC contracts by default; under +# icpx, after the project-wide -fp-model=precise, which would allow it). +set_source_files_properties(src/foldquant.cpp src/foldquant_ref.cpp PROPERTIES COMPILE_OPTIONS "-ffp-contract=off") if(VLA_SPM) target_include_directories(vla_core PRIVATE ${sentencepiece_SOURCE_DIR}/src) target_link_libraries(vla_core PRIVATE sentencepiece-static ${VLA_SPM_PROTOBUF}) @@ -268,8 +300,36 @@ if(GGML_CUDA) target_compile_definitions(vla_core PUBLIC VLA_BITVLA_CUDA_KERNELS) target_include_directories(vla_core PUBLIC ${CUDAToolkit_INCLUDE_DIRS}) + # FoldQuant kernels: their own archive, device-linked on its own like bitvla + # above, because they compile with -fmad=false (bit-exact against the CPU + # reference in src/foldquant_ref.h) and no other archive does. + add_library(vla_fq_kernels STATIC + src/kernels/foldquant/fq_prologue.cu + src/kernels/foldquant/fq_gemm_i8.cu + src/kernels/foldquant/fq_gemm_mma.cu + ) + set_target_properties(vla_fq_kernels PROPERTIES + CUDA_SEPARABLE_COMPILATION ON + CUDA_RESOLVE_DEVICE_SYMBOLS ON + POSITION_INDEPENDENT_CODE ON + ) + target_include_directories(vla_fq_kernels PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) + target_compile_features(vla_fq_kernels PRIVATE cxx_std_17) + target_compile_options(vla_fq_kernels PRIVATE + $<$:-O3 -Xptxas=-O3 -fmad=false> + ) + target_link_libraries(vla_fq_kernels PUBLIC CUDA::cudart) + # ggml-cuda link: the hook resolves ggml_cuda_ext_forward out of it. - add_library(vla_cuda_ops STATIC src/cuda/vla_cuda_bf16.cu) + add_library(vla_cuda_ops STATIC + src/cuda/vla_cuda_ext.cu + src/cuda/vla_cuda_bf16.cu + src/cuda/vla_cuda_foldquant.cu + ) + # The staging shim in vla_cuda_foldquant.cu runs the CPU reference on the + # host; it must round exactly like vla_core's copy. + set_source_files_properties(src/cuda/vla_cuda_foldquant.cu PROPERTIES + COMPILE_OPTIONS "-Xcompiler=-ffp-contract=off") target_include_directories(vla_cuda_ops PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src ${llama_SOURCE_DIR}/ggml/include @@ -283,7 +343,7 @@ if(GGML_CUDA) target_compile_options(vla_cuda_ops PRIVATE $<$:-O3 -Xptxas=-O3> ) - target_link_libraries(vla_cuda_ops PUBLIC CUDA::cublas CUDA::cudart ggml-cuda) + target_link_libraries(vla_cuda_ops PUBLIC CUDA::cublas CUDA::cudart ggml-cuda vla_fq_kernels) target_link_libraries(vla_core PRIVATE vla_cuda_ops) endif() @@ -296,6 +356,29 @@ if(GGML_SYCL AND NOT GGML_CUDA) "See docs/backend/sycl.md.") endif() target_compile_definitions(vla_core PUBLIC GGML_USE_SYCL) + + # FoldQuant kernels for Intel GPUs, behind the hook the patch above adds. Built + # without FMA contraction, and with the device JIT told to round division and + # sqrt correctly (the GPU driver approximates both otherwise), so the codes + # match the CPU reference bit for bit. The JIT option is a link option: it + # rides in the executable's device image. One image per kernel, so a device + # JIT-compiles only what it launches (not the XMX kernel on a device without + # XMX: the OpenCL CPU device's compiler has crashed on the whole module). + target_sources(vla_core PRIVATE src/sycl/vla_sycl_foldquant.cpp) + set_source_files_properties(src/sycl/vla_sycl_foldquant.cpp PROPERTIES + COMPILE_OPTIONS "-fsycl;-ffp-contract=off") + target_link_options(vla_core PUBLIC -fsycl -fsycl-device-code-split=per_kernel + "SHELL:-Xsycl-target-backend=spir64 \"-cl-fp32-correctly-rounded-divide-sqrt\"") + target_compile_definitions(vla_core PRIVATE VLA_FQ_SYCL) + # Prefill-sized FoldQuant GEMMs go to oneDNN's int8 matmul when ggml-sycl + # was built with it (GGML_SYCL_DNN, the default). + if(GGML_SYCL_DNN) + find_package(DNNL QUIET) + if(DNNL_FOUND) + target_link_libraries(vla_core PRIVATE DNNL::dnnl) + target_compile_definitions(vla_core PRIVATE VLA_FQ_SYCL_DNNL) + endif() + endif() endif() # Precedence matches the ladder in src/backend.h. @@ -307,6 +390,8 @@ if(GGML_OPENVINO AND NOT GGML_CUDA AND NOT GGML_SYCL AND NOT GGML_METAL) # Only tells the archs which branch of the backend.h ladder to compile; the # toolkit was located before the fetch above. target_compile_definitions(vla_core PUBLIC GGML_USE_OPENVINO) + # FoldQuant nodes run natively through the translator compiled into ggml-openvino. + target_compile_definitions(vla_core PRIVATE VLA_FQ_OPENVINO) endif() # Hexagon and OpenCL reject some ops the archs build; the wrapper in @@ -461,7 +546,7 @@ if(VLA_BUILD_SERVER) endif() # The CUDA targets were the only first-party code compiled without warnings. if(GGML_CUDA) - list(APPEND VLA_FIRST_PARTY_TARGETS bitvla_cuda_kernels vla_cuda_ops) + list(APPEND VLA_FIRST_PARTY_TARGETS bitvla_cuda_kernels vla_cuda_ops vla_fq_kernels) endif() option(VLA_WERROR "Treat warnings in first-party code as errors" OFF) diff --git a/README.md b/README.md index 9ea731d..2f73c8e 100644 --- a/README.md +++ b/README.md @@ -225,6 +225,7 @@ Wiring, recording, training and queue sizing are in the | [docs/USAGE.md](docs/USAGE.md) | `vla-cli`, `vla-server`, `vla-bench`: prompt tokenization, `-hf` tags, runtime flags, environment variables | | [docs/EVAL.md](docs/EVAL.md) | Installing LIBERO and SimplerEnv, running the eval clients against `vla-server` | | [docs/MODELS.md](docs/MODELS.md) | Converting safetensors checkpoints to GGUF, quantizing to Q8_0/Q4_0 | +| [docs/QUANTIZATION.md](docs/QUANTIZATION.md) | FoldQuant W8A8 / W4A4: the GGUF contract, the integer arithmetic, converting a FoldQuantVLA quantized model | | [docs/DOCKER.md](docs/DOCKER.md) | Building and running the eval in containers | | [docs/ARCHITECTURE.md](docs/ARCHITECTURE.md) | Engine design: layers, the prediction path, backends, adding an architecture | | [docs/backend/](docs/backend) | Per-backend build and run notes: [SYCL](docs/backend/sycl.md), [OpenVINO](docs/backend/ov.md), [Metal](docs/backend/metal.md), [Hexagon](docs/backend/hexagon.md), [Hexagon on Windows](docs/backend/hexagon-windows.md), [WSL2](docs/backend/wsl.md) | diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 4ce8426..7f7939b 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -30,6 +30,8 @@ source is the detail. - `src/serving/` - `vla-server` (ZeroMQ + protobuf, action prediction), `vlm-server` (chat), `vla-cli` (one-shot inference) and `vla-bench` (timing). - `src/kernels/bitvla/` - custom 1.58-bit ternary CUDA kernels for BitVLA. +- `src/foldquant.h`, `src/kernels/foldquant/`, `src/cuda/` - FoldQuant INT8 linears + and the in-tree CUDA kernels behind the ggml extension hook. ## The prediction path @@ -70,9 +72,17 @@ llama.cpp is fetched by CMake `FetchContent` and pinned by `VLA_LLAMA_TAG` in (`--weight-dtype` overrides it), and a GGUF can be repacked to Q8_0/Q4_0 with `scripts/quantize_gguf.py`. The loader keeps packed weights packed. On CPU and CUDA, ggml quantizes the -activations to 8 bits and runs int8 dot products on the blocks. CPU thread count -scales to the machine core count; the GPU backends run the towers and the -transformer on the device. +activations to 8 bits and runs int8 dot products on the blocks. + +A FoldQuant GGUF (see [QUANTIZATION.md](QUANTIZATION.md)) +carries INT8 or INT4 codes plus sidecar scales instead. `src/foldquant.h` declares those +sites, `src/layers/fq_linear.h` turns each into two `GGML_OP_CUSTOM` nodes, the +CPU backend runs the reference in `src/foldquant_ref.cpp`, and on CUDA the +`src/kernels/foldquant/` integer kernels claim the same nodes through the ggml +extension hook (`src/cuda/`), and on OpenVINO `src/openvino/foldquant_ov.cpp` +translates them into OpenVINO ops. Every other backend reads the sites back as +float weights through `WeightLoader::as_float` and runs the arch's float path. CPU thread count scales to the machine core count; +the GPU backends run the towers and the transformer on the device. ## Adding an architecture diff --git a/docs/MODELS.md b/docs/MODELS.md index f1b8787..75b8b3c 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -43,3 +43,29 @@ compute, not BF16. a bigger cut (`--type` also takes `Q4_1`, `Q5_0`, `Q5_1`). Embeddings, the output head, norms and the action expert stay float; pass `--vision` to pack the vision tower too (smaller, but more accuracy loss). + +### FoldQuant W8A8 / W4A4 + +The stock repack keeps the action head float and rounds the LM weights block by +block with no calibration. A FoldQuant GGUF ships the language backbone and the +action head as INT8 or INT4 codes in a Hadamard-rotated, SmoothQuant-folded frame +with per-row scales, calibrated by FoldQuantVLA; vla.cpp quantizes the activations +per token (to 4 bits for W4A4) and runs the GEMMs on the integer tensor cores. It loads like any other checkpoint on every backend: CUDA and SYCL (Intel GPUs) run +integer kernels, CPU the exact reference and OpenVINO (CPU and GPU) the same +arithmetic as OpenVINO ops with INT8/INT4 weight constants; the others read the sites back as float +weights with FoldQuant's rounding (weight-only quantization, bf16-GGUF speed); the format +and the arithmetic are in [QUANTIZATION.md](QUANTIZATION.md). + +A calibrated arm saved as a quantized model by +[FoldQuantVLA](https://github.com/VinRobotics/FoldQuantVLA) converts to that file with +no calibration and nothing re-rounded (GR00T N1.5 / N1.6 / N1.7 and π0.5): + +```bash +python scripts/convert_quantized_model_to_gguf.py --quantized-model --out model-fq.gguf +./build/vla-server model-fq.gguf --bind tcp://*:5556 + +# uncalibrated stand-in for bring-up and benchmarks (rotation + per-row INT8, no SmoothQuant) +python scripts/foldquant_fake_export.py --in model-bf16.gguf --out model-fq.gguf +python scripts/inspect_gguf_quant.py model-fq.gguf +./build/vla-bench --ckpt model-fq.gguf --images 1 --size 256 --tokens 16 +``` diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md new file mode 100644 index 0000000..027f7d3 --- /dev/null +++ b/docs/QUANTIZATION.md @@ -0,0 +1,314 @@ +# FoldQuant GGUF: the file contract + +vla.cpp runs two kinds of quantized checkpoint: + +| | Stock repack | FoldQuant | +|---|---|---| +| Producer | `scripts/quantize_gguf.py` | `scripts/convert_quantized_model_to_gguf.py` from a FoldQuantVLA quantized model (calibrated), `scripts/foldquant_fake_export.py` (uncalibrated) | +| Weights | ggml `Q8_0` / `Q4_0` blocks (block-32 absmax) | INT8 or INT4 codes, per-output-row scale, block-Hadamard-rotated frame, SmoothQuant folded in | +| Activations | float | dynamic per-token INT8 (INT4 in phase 3) | +| Executed by | `ggml_mul_mat` dequantizing at compute | in-tree integer kernels (`src/kernels/foldquant/`) or the CPU reference | +| Backends | all | integer kernels on CUDA, exact reference on CPU, OpenVINO ops on the OpenVINO CPU/GPU plugins; every other backend reads the sites back as float weights (see [Other backends](#other-backends)) | + +This page is the canonical description of the FoldQuant file and of the +arithmetic the runtime performs on it. The converters and vla.cpp's loader +and kernels are all written against it; `scripts/inspect_gguf_quant.py` +checks a file against it. + +## The file + +A FoldQuant GGUF is the file the architecture's converter +(`scripts/convert__to_gguf.py`) would write - same tensor names, same +`.*` keys, BF16 elsewhere - except that each quantized **site** replaces +its float weight with: + +| Tensor | Type, ggml `ne` | Meaning | +|---|---|---| +| `.weight` | `I8`, `(K, N)`; INT4: `(K/2, N)` nibble-packed | integer codes, output row `n` contiguous along `k`, in the rotated (and SmoothQuant-folded) frame | +| `.wscale` | `F32 (N)` | per-output-row weight scale | +| `.ascale` | `F32 (K)`, optional | the static SmoothQuant vector `s` the activation is **divided** by; action-module sites only | +| `.bias` | unchanged | | +| norm gammas | unchanged names, **folded** values | the LLM's `attn_norm` / `ffn_norm` already carry SmoothQuant | + +`ascale` is the shipped static vector. The dynamic per-token activation +scale is computed at runtime and never stored. + +INT4 nibbles: byte `b` of a row holds column `2b` in the low nibble and +`2b+1` in the high nibble, two's complement, values in `[-7, 7]`. + +Sites of GR00T N1.5 / N1.6 / N1.7 (KV prefix `gr00t_n1_5` / `gr00t_n1_6` / `gr00t_n1_7`): + +- LLM: `vlm.blk.{i}.{attn_q, attn_k, attn_v, attn_o, ffn_gate, ffn_up, ffn_down}` +- DiT: `aex.dit.{i}.{attn_q, attn_k, attn_v, attn_o, ff0, ff2}` + +Sites of pi0.5 (KV prefix `pi05`): + +- PaliGemma prefix (LLM recipe): `vlm.blk.{i}.{attn_q, attn_k, attn_v, attn_o, + ffn_gate, ffn_up, ffn_down}`. The RMSNorm and its folded gamma ride in the + q/k/v and gate/up act nodes. vla.cpp loads Gemma norms as `1 + w`, so the + exporter writes `attn_norm` / `ffn_norm` as F32 holding the folded gamma + minus one. +- Gemma action expert (action recipe): `aex.blk.{i}.{attn_q, attn_k, attn_v, + attn_o, ffn_gate, ffn_up, ffn_down}`, each with an `ascale` (its input comes + out of AdaRMS, so there is no gamma to fold into); q/k/v and gate/up share + one vector per group. The expert's residuals are gated, so its o / down + GEMMs do not take the fused residual. + +q/k/v (and gate/up) are separate tensors with their own `wscale`. Projections +that read the same input share one input transform, so where the loader fuses +them the codes and `wscale` concatenate along `N` and any `ascale` must be +identical across the group - the exporter writes the same vector under each +name, the loader asserts equality and keeps one. The fused groups are the DiT's +`Wqkv` on self-attention blocks and `Wkv` on cross-attention blocks; on a +cross block `attn_q` reads the hidden state while `attn_k`/`attn_v` read the +VL encoder, so `q` carries its own `ascale` (and a different `K`) there. +`scripts/inspect_gguf_quant.py` infers the split from `K`. ViT, VLSA, adaLN, the state/action +encoders and decoders stay float in every scheme. + +`K` and `N` must be multiples of 64 at every site. A site whose shape does not +qualify must be left float by the exporter. + +## Metadata + +All keys are prefixed `.quant.` (e.g. `gr00t_n1_7.quant.method`). + +| Key | Type | Meaning | +|---|---|---| +| `method` | str | `"foldquant"` - the presence test | +| `applied_at` | str | producer (`"foldquant"`, `"foldquant_fake_export"`) | +| `scheme_llm`, `scheme_action` | str | the producer's scheme keys (FoldQuantVLA `w8a8_sr`, `w4a4_shg`, ...), for logs and provenance | +| `llm_weight_bits`, `llm_act_bits` | u32 | 8 or 4 | +| `llm_rot_block_size` | u32 | nominal Hadamard block (64) | +| `action_weight_bits`, `action_act_bits`, `action_rot_block_size` | u32 | same for the action module | +| `action_fold_order` | str | `"before"`: divide by `ascale` before the butterfly; `"after"`: after | +| `act_clip_ratio` | f32 | activation clip, INT4 activations only (LLM) | +| `site_bits` | str | LLM per-site overrides, e.g. `"o:8,down:8"`; keys `qkv \| o \| gateup \| down`; absent keys inherit the module widths | +| `provenance` | str | free text: the quantized model it came from, its base checkpoint and digest | + +The rotation block is nominal: both sides narrow it per site to the largest +power of two that divides `K` (`rotation_block_for`), 1 meaning no rotation. + +## The arithmetic + +Per site, on a token row `x[K]` in fp32, in this order: + +1. `[gamma]` RMSNorm with the folded gamma: `y = (x * rstd) * gamma`, + `rstd = 1 / sqrt(mean(x^2) + eps)`. LLM q/k/v and gate/up sites; the norm is + fused into the activation node. `attn_o` / `ffn_down` take the raw residual + stream and skip this step; DiT sites take the adaLN / LayerNorm output. +2. `[fold_order = before, ascale]` `y = y / ascale` +3. `[rot_block > 1]` in-place block butterfly on every `rot_block` elements: + stages `h = 1, 2, 4, ...`, pairs `(i, i+h)` become `(a+b, a-b)`, then + `y *= 1/sqrt(rot_block)` (a host-computed float constant). +4. `[fold_order = after, ascale]` `y = y / ascale` +5. `scale = max(clip * amax(|y|) / qmax, 1e-12)`, `inv = 1 / scale`, + `q = clamp(rint(y * inv), -qmax, qmax)`, `qmax = 127` (INT8) or `7` (INT4), + `rint` rounding half to even. The reciprocal multiply is what FoldQuant's + TensorRT kernels do (`rmsnorm_per_row_quant_cuda.cu`), so engine and runtime + round alike; PyTorch's `x / scale` emulation can differ by one code at a + rounding boundary. +6. `acc[n] = sum_k q[k] * w[n][k]` in int32; `out[n] = ((float) acc * scale) * wscale[n] (+ bias[n])`. + +Weights were prepared offline as `W' = W . H_block^T` (plus the SmoothQuant +fold) and rounded per output row with `wscale = amax_row / qmax`, so `W' y'` +equals `W x` up to quantization. + +`tests/foldquant_gemm_check` (CUDA builds, not ctest) times the prologue and +GEMM at GR00T shapes against a cuBLAS BF16 GEMM of the same shape, and +`--stress` re-runs the FFN shapes hundreds of times hashing the outputs, which +is how a kernel race would show. Do not run any of these, or `vla_predict_check`, +while a build relinks `libvla_core.so` / `libggml-cuda.so`: a process that has +the old library mapped executes a mix of old and new code and its output is +garbage that looks exactly like a race. + +`src/foldquant_ref.cpp` is this list written out in the kernels' operation order +(per-thread strided partial sums for `mean(x^2)`, a fixed reduction tree, IEEE +sqrt and division, no FMA contraction). It is the one translation unit built +with `-ffp-contract=off`; keep the bodies there rather than in the header, since +an inline copy compiled into a model file picks up the default contraction and +the linker keeps whichever copy it likes. The CPU backend runs it as is; the CUDA +kernels (`src/kernels/foldquant/`, compiled with `-fmad=false`) reproduce it bit +for bit - `tests/test_foldquant_cuda_op.cpp` checks that on isolated nodes and +`VLA_FQ_CHECK=1` re-checks every node inside a real model run. `scripts/foldquant_ref.py` +is the same arithmetic in numpy; `tests/py/test_foldquant_ref.py` pins it to the +C++ test's golden checksum. + +## In the graph + +Each site is two `GGML_OP_CUSTOM` nodes (`src/layers/fq_linear.h`): + +- `fq_act(x[, gamma][, ascale]) -> I8 [K_pack + 16, T]`: per row, the codes + (`K` bytes, or `K/2` nibble-packed) followed by the fp32 per-token scale at + byte `K_pack`. Sources are packed without holes (the spec's `has_gamma` / + `has_ascale` say which follow `x`). It is an ordinary gallocr intermediate, + shared by every projection that reads the same input (q/k/v, gate/up). +- `fq_gemm(w, blob, wscale[, bias][, residual]) -> F32 [N, T]`. The fifth + source is the F32 tensor the model would add right after the GEMM (o_proj and + down/ff2 residuals); the epilogue adds it, one float add, so the result is + what `ggml_add` would give. When the spec carries a head layout + (`fq_set_heads`: DiT q/k/v, cross-attention k/v, LLM v) the epilogue writes + each projection straight in the layout the attention reads, `[hd, T, heads]` + for Q/K or `[T, hd, heads]` for V, and the model takes views instead of + permute copies; the numbers are the same, only their addresses move. + +The CPU backend executes the custom function. On CUDA the same nodes are +claimed by the extension hook (`src/cuda/vla_cuda_foldquant.cu`, registered by +`foldquant_check_backend` at load) through the magic word in the node's +userdata; a node that violates the contract is declined, and ggml then aborts +on the unsupported op rather than computing something else. + +### OpenVINO + +On ggml's OpenVINO backend the two nodes are translated into OpenVINO ops +(`src/openvino/foldquant_ov.cpp`, compiled into the backend; the hook that +registers it for `GGML_OP_CUSTOM` is hunk 14 of `scripts/patch_ggml_openvino.py`). +`fq_act` becomes the reference's arithmetic in OpenVINO ops (RMSNorm with the +folded gamma, the ascale divide, the block rotation as a MatMul with the +normalised Hadamard matrix, `scale = max(clip * amax / qmax, 1e-12)`, +`clamp(round_half_even(y * (1 / scale)))`) and carries the codes and the +per-token scale as floats to `fq_gemm`, which multiplies them with the weight +kept as an `i8` or `i4` constant (the W4 nibble bytes are reinterpreted in place; +OpenVINO's `i4` has the same low-nibble-first order) dequantized by `wscale` in +the decompression pattern the plugins keep compressed. The head-laid-out +epilogue is off there (`FqModuleSpec::no_heads`), since a translated graph has +no raw layout for the views to read. + +That runs on the CPU and GPU plugins (`GGML_OPENVINO_DEVICE=CPU|GPU`). The NPU +compiler accepts no such graph, so on the NPU the file uses dequant mode below, +as `VLA_FQ_DEQUANT=1` does anywhere. `tests/test_foldquant_ov_op.cpp` runs one +site through the backend against the CPU reference: every code agrees and the +outputs match to under 1e-6 relative, on an Intel CPU and an Arc iGPU. A whole +model follows the CUDA integer path to 1.00000 action cosine at W8A8; at W4A4 a +code that sits on a rounding tie in one place and not the other (the reductions +run in a different float order) flips by one, and with 15 activation levels +those flips compound through the layers and the denoise steps to about 0.999. + +### SYCL + +On ggml's SYCL backend (Intel GPUs) the nodes are claimed through an extension +hook like CUDA's (`scripts/patch_ggml_sycl_ext_hook.py` adds it to ggml-sycl; +`src/sycl/vla_sycl_foldquant.cpp` registers the kernels at load). `fq_act` +follows the CPU reference's reduction tree: lane `l` of a 32-wide sub-group owns +the 64-element chunks `l, l+32, ...` of a row, and the sum of squares and the +amax are xor butterflies over the sub-group. `fq_gemm` runs oneDNN's int8 matmul +(INT4 weights as oneDNN `s4`, INT4 activations unpacked to `s8`) into an int32 +buffer, then the reference's epilogue. Built without oneDNN (`GGML_SYCL_DNN=OFF`), +or with `VLA_FQ_SYCL_GEMM=native`, it uses its own GEMV for up to 32 tokens and +an XMX `joint_matrix` kernel above that; both are much slower. + +The integer sums are exact in any order, so the output is bit-identical to the CPU +reference as long as the float steps are: the source is built with +`-ffp-contract=off`, and the device image carries +`-cl-fp32-correctly-rounded-divide-sqrt` for the driver's JIT, which otherwise +approximates division and sqrt and moves codes across rounding ties. +`tests/test_foldquant_sycl_op.cpp` checks every activation byte and every output +bit against the reference, for every bit width, on every GEMM path, at +production shapes up to K = 16384. With a GPU it goes through ggml's SYCL +backend; without one (CI, or `ONEAPI_DEVICE_SELECTOR=opencl:cpu`) it hands the +nodes straight to the extension hook on the OpenCL CPU device, since ggml-sycl +will not start on a CPU; the GEMMs then run on vla.cpp's own kernels, because +oneDNN's int8 matmul on a CPU without VNNI saturates 16-bit partial sums. A +whole π0.5 model follows the CUDA integer path to 0.99999 action cosine at W8A8, 0.9997 at W4A4 with INT8 o/down and 0.9995 +at W4A4 (FoldQuantVLA's calibrated LIBERO checkpoints). Both GPUs are exact to the same reference, so the gap comes from the +float layers between the sites (the bf16 model's actions differ by up to 7e-4 +between them), which land some codes on the other side of a rounding tie; with +15 activation levels those flips compound, as on OpenVINO. + +### Other backends + +Metal, Vulkan, Hexagon and OpenCL (and the OpenVINO NPU) have no +implementation of the two custom nodes, and vla.cpp drives a single backend with +no per-op fallback. +There `foldquant_check_backend` switches the file to dequant mode: every site is +registered with the loader as a float GEMM weight, rebuilt at upload in the +resident type (`--weight-dtype`) from its codes, `wscale` and `ascale`, and the +arch takes its stock float path. The activation path above is +`x' = R(x / a)` (`fold_order = before`) or `R(x) / a` (`after`), with `R` the +block-normalised Sylvester-Hadamard butterfly, which is symmetric and +orthonormal; so `W_deq x' = (W_deq R diag(1/a)) x` or `(W_deq diag(1/a) R) x`, and +row `n` of the float weight is `R(w_n) / a` or `R(w_n / a)` (`fq_dequant_rows`). +The LLM sites' SmoothQuant vector is already folded into the norm gains the file +carries, which the float path reads as its norm weights. + +What that runs is weight-only quantization: the weights keep FoldQuant's +rounding, the activations stay float (no per-token quantization, no INT4 clip), +so the actions are close to, not bit-identical with, the integer path, and +memory and speed are those of the bf16 GGUF. `VLA_FQ_DEQUANT=1` selects the same +mode on CUDA or CPU (on a CPU it is much faster than the exact reference). +`tests/test_foldquant_dequant.cpp` checks the rebuilt weights against the dense +product. + +Environment switches: `VLA_FQ_CHECK=1` (CUDA and SYCL) recomputes every node +with the CPU reference after its kernel and reports mismatches; `VLA_FQ_CPU_REF=1` runs the CPU reference on host copies +of every node (a byte-exact A/B against the kernels; it disables ggml's CUDA +graphs, whose capture cannot contain the host round trip); `VLA_FQ_TRACE=1` +prints each node's shape once per graph build. A/B switches for the graph-level +optimisations, all bit-identical either way: `VLA_FQ_NO_FUSE=1` keeps the +residual add as a separate node, `VLA_FQ_NO_HEADS=1` keeps the permute copies, +`VLA_FQ_PREFETCH_MB=` makes each GEMM prefetch that much of the next site's +weights into L2 (opt-in; measured slower on Orin). + +## Producing a file + +- The family converters (`scripts/convert__to_gguf.py`) expose + `convert(ckpt, out, *, writer_factory, ...)`: a writer factory from + `scripts/gguf_quant_writer.py` turns their output into a FoldQuant file. + `convert_pi05_to_gguf.py` also takes an OpenPI-converted checkpoint with + `--config-json` (the lerobot policy fields) and an OpenPI `norm_stats.json`. +- From a quantized model (a quantized policy saved as data: the calibration result and + every site's integer codes, next to the base checkpoint), with no calibration and + nothing re-rounded: + `python scripts/convert_quantized_model_to_gguf.py --quantized-model --out model.gguf`. + FoldQuantVLA writes one in two formats, and the converter takes both (told apart by the + manifest's `format`, not its file name): + - FoldQuantVLA's quantized checkpoint (`foldquant.quantized_checkpoint`, format + `foldquant-quantized-checkpoint` v1): the base checkpoint in which every quantized + projection's `.weight` is replaced by `.qweight` (int8 `(N, K)`, or int4 nibble-packed + `(N, K/2)` uint8, low nibble = even column) and `.weight_scale` (fp32 per row), the + SmoothQuant vectors as `foldquant..sq.` tensors, and `foldquant_quant.json` + naming, per site, the checkpoint keys it covers in row order. The directory is its own + base checkpoint (`--ckpt` is refused). The family converter reads it through a base view + in which each missing `.weight` is `qweight x weight_scale`; sites are then written as + codes. A W4A4 DiT also stores its adaLN projections as INT4; vla.cpp keeps adaLN in float, + so those are written from the dequantized weight, with the scale rounded to bf16 first as + the TensorRT `AdaLNModInt4` plugin bakes it. Every quantized projection must end up as a + GGUF site or as that adaLN, or the conversion fails. + - FoldQuantVLA's earlier fake-quant state (`foldquant.fakequant`, format + `foldquant-quant-state`: `foldquant_quant.json` + `quant_state.safetensors` beside the + base checkpoint). + + For both FoldQuantVLA formats the LLM norm gains are folded as FoldQuant's emitter folds + them (`(w [+ 1 for Gemma]) / s` in fp32 on the gain in the export dtype, rounded once), and + `--check-onnx /onnx` byte-compares every site and gain against the plugin ONNX graphs + the TensorRT engines were built from and fails on any difference. + + The recorded codes and row scales go into the GGUF as they are (FoldQuant's INT4 layout + is this page's) and each action site ships its SmoothQuant vector as `.ascale`. + + | family | LLM schemes | action schemes | checked | + |---|---|---|---| + | GR00T N1.5 / N1.6 / N1.7 | `w8a8_sr`, `w4a4_srg` (+ `site_bits`) | `w8a8_sh`, `w4a4_shg` | fake-quant states `n15_fq`, `n16_fq_clean`, `fq_w8v2`, `fq_w4v2` (LIBERO) and N1.7 SO101 `w8a8` / `w4a4` / `w4a4_od8`: every site and gain byte-identical to the TensorRT graphs; N1.7 SO101 `w8a8` run on the real arm. Quantized checkpoints: unit-tested | + | pi0.5 | `w8a8_sr`, `w4a4_srg` (+ `site_bits`, activation clip) | `w8a8_sh`, `w4a4_sh`, `w4a4_shg` | FoldQuantVLA `pi05_fq` (LIBERO, `w8a8_sr` + `w4a4_shg`) byte-identical to its TensorRT graphs; action cosine median 0.99991 vs 0.99986 for the engines | + | any | | `*_sr` (dense learned rotation) | refused: vla.cpp runs the fixed butterfly only | + + In the earlier FoldQuantVLA state, GPTQ (INT4) DiT sites are recorded by name and + round-to-nearest (INT8) DiT packs by call order (`dit.rtn/{n}`); the converter walks the latter in the emitter's + per-block order (cross: q, o, kv, ffn0, ffn2; self: qkv, o, ffn0, ffn2), checks every + shape, and `--check-onnx` confirms the mapping byte for byte. The DiT's adaLN projections + stay float in vla.cpp, so that format's adaLN packs are not needed. A Hub-style state without the + base checkpoint files takes `--ckpt`. Action `sq_fold_order` must be `before` (the only order the + runtime implements); a learned per-site activation clip has no metadata slot yet. + +- Uncalibrated, for kernel bring-up and benchmarks: + `python scripts/foldquant_fake_export.py --in n17-bf16.gguf --out n17-fq.gguf`. +- Check: `python scripts/inspect_gguf_quant.py n17-fq.gguf` (exit 1 on a violation). + +## Phases + +1. W8A8, butterfly rotation, GR00T N1.6 / N1.7 (this page). +2. W4A8: INT4 codes (`*_weight_bits = 4`, `ne0 = K/2`), INT8 activations; the + same file layout otherwise. The CUDA GEMM runs the CPU reference for these + sites until the nibble-unpacking kernel lands. +3. W4A4: INT4 activations with `act_clip_ratio`. +4. Other families: GR00T N1.5 and pi0.5 are wired; pi0, SmolVLA and Evo-1 are not + (per-module validation taps and dense rotations are also still open). diff --git a/docs/backend/ov.md b/docs/backend/ov.md index 2393349..e702ffc 100644 --- a/docs/backend/ov.md +++ b/docs/backend/ov.md @@ -287,6 +287,42 @@ graph at a stage, which makes that stage the terminal node, or set of the subtlest fixes above were found by bisecting that way, comparing each stage against a CPU-backend reference; neither logs anything when it goes wrong. +## FoldQuant (W8A8 / W4A4) checkpoints + +A FoldQuant GGUF runs on the CPU and GPU plugins through +`src/openvino/foldquant_ov.cpp`, which translates FoldQuant's two custom nodes +into OpenVINO ops with the weights kept as `i8` / `i4` constants (hunk 14 of +`scripts/patch_ggml_openvino.py` registers it). The NPU, and `VLA_FQ_DEQUANT=1`, +read the sites back as float weights instead. See +[QUANTIZATION.md](../QUANTIZATION.md#openvino) for the arithmetic and accuracy. + +π0.5 LIBERO, 2 views at 224 px, 48 tokens, Intel Core Ultra X7 358H with an Arc +B390 iGPU, OpenVINO 2026.4, `vla-bench` p50 over 10 calls after 3 warmups: + +| Device | Checkpoint | p50 ms | Peak RSS | +|---|---|--:|--:| +| CPU | bf16 | 6624 | 16.5 GB | +| CPU | FoldQuant W4A4 | 3888 | 6.4 GB | +| CPU | FoldQuant W4A4, o/down INT8 | 3807 | 7.1 GB | +| CPU | FoldQuant W8A8 | 3653 | 8.6 GB | +| GPU | bf16 | 568 | 11.4 GB | +| GPU | FoldQuant W4A4 | 656 | 3.4 GB | +| GPU | FoldQuant W4A4, o/down INT8 | 652 | 3.7 GB | +| GPU | FoldQuant W8A8 | 655 | 4.7 GB | + +The GPU runs the per-site activation ops (reduction, rounding, the rotation +MatMul) on top of the GEMMs and comes out slower than bf16 there, at a third of +the memory. On an Intel GPU the SYCL backend is the faster FoldQuant path: the +same W4A4 file runs in 291 ms on this iGPU, against 326 ms for bf16 on SYCL +([sycl.md](sycl.md#foldquant-w8a8--w4a4-checkpoints)). + +The FoldQuant files are FoldQuantVLA quantized checkpoints converted with +`scripts/convert_quantized_model_to_gguf.py`: W4A4 and W4A4 with INT8 o/down are +[vrfai/pi05-libero-w4a4](https://huggingface.co/vrfai/pi05-libero-w4a4) and +[vrfai/pi05-libero-w4a4-res8](https://huggingface.co/vrfai/pi05-libero-w4a4-res8); +W8A8 is FoldQuantVLA's W8A8 arm (`--llm-scheme w8a8_sr --expert-scheme w8a8_sh`) +calibrated on the same 128 LIBERO frames, seed 0. + ## Known issues **π0 needs F32 on the GPU, and gets it by default.** The GPU plugin computes in diff --git a/docs/backend/sycl.md b/docs/backend/sycl.md index 01adcd5..65b97ec 100644 --- a/docs/backend/sycl.md +++ b/docs/backend/sycl.md @@ -126,6 +126,48 @@ BitVLA is the one exception: it pins its ggml graph to the CPU backend by design and offloads its LM through separate hand-written CUDA kernels, so a SYCL build leaves it on the CPU. There is no SYCL port of those kernels. +## FoldQuant (W8A8 / W4A4) checkpoints + +A FoldQuant GGUF runs its integer path on SYCL: `src/sycl/vla_sycl_foldquant.cpp` +claims FoldQuant's two custom nodes through an extension hook that +`scripts/patch_ggml_sycl_ext_hook.py` adds to ggml-sycl at configure time. The +activation quantization runs in its own kernels and the GEMMs in oneDNN's int8 +matmul, bit-identical to the CPU reference. `VLA_FQ_DEQUANT=1` reads the sites back +as float weights instead. See [QUANTIZATION.md](../QUANTIZATION.md#sycl) for the +arithmetic and accuracy. + +π0.5 LIBERO, 2 views at 224 px, 48 tokens, Intel Core Ultra X7 358H with an Arc +B390 iGPU, oneAPI 2026.1, `vla-bench` p50 over 10 calls after 3 warmups: + +| Checkpoint | GGUF | p50 ms | +|---|--:|--:| +| bf16 | 6.71 GB | 326 | +| FoldQuant W4A4, `VLA_FQ_DEQUANT=1` | 3.27 GB | 327 | +| FoldQuant W4A4 | 3.27 GB | **291** | +| FoldQuant W4A4, o/down INT8 | 3.61 GB | **290** | +| FoldQuant W8A8 | 4.42 GB | **305** | + +Vision (79 ms) stays bf16 in every row. W4A4 halves the weights and is 11% faster +than bf16, W8A8 6%. + +The FoldQuant files are FoldQuantVLA quantized checkpoints converted with +`scripts/convert_quantized_model_to_gguf.py`: W4A4 and W4A4 with INT8 o/down are +[vrfai/pi05-libero-w4a4](https://huggingface.co/vrfai/pi05-libero-w4a4) and +[vrfai/pi05-libero-w4a4-res8](https://huggingface.co/vrfai/pi05-libero-w4a4-res8); +W8A8 is FoldQuantVLA's W8A8 arm (`--llm-scheme w8a8_sr --expert-scheme w8a8_sh`) +calibrated on the same 128 LIBERO frames, seed 0. + +## Known issue: the Level Zero loader on Ubuntu 24.04 + +oneAPI 2026.1's Level Zero adapters crashed (`sycl-ls` segfaults) against Ubuntu +24.04's `libze1` 1.16 on the machine above, most likely a loader too old for them +(a newer loader from Intel's GPU packages is untested). Pointing the runtime at +the OpenCL adapter works, and the numbers above were measured that way: + +```bash +export UR_ADAPTERS_FORCE_LOAD=/opt/intel/oneapi/compiler/2026.1/lib/libur_adapter_opencl.so.0 +``` + ## Known issue: the SYCL VMM pool and oneDNN ggml-sycl's VMM pool hands out virtual-memory-backed pointers that oneDNN cannot diff --git a/scripts/convert_gr00t_n1_5_to_gguf.py b/scripts/convert_gr00t_n1_5_to_gguf.py index 70a846d..36f8ea4 100644 --- a/scripts/convert_gr00t_n1_5_to_gguf.py +++ b/scripts/convert_gr00t_n1_5_to_gguf.py @@ -130,28 +130,14 @@ def _write_lerobot_stats(ckpt: Path, out_path: Path, emb_key: str = "new_embodim out_path.write_text(json.dumps(blob, indent=2)) print(f"wrote {out_path} (embodiment {emb_key!r}: state[{len(s_min)}] + action[{len(a_min)}] min/max)") -def main() -> int: - ap = arg_parser( - ARCH, - "GR00T-N1.5-3B checkpoint dir. Two layouts are auto-detected: " - "(a) NVIDIA Isaac snapshot (sharded safetensors, un-prefixed tensor names, " - "config.json model_type=gr00t_n1_5, experiment_cfg/metadata.json); " - "(b) lerobot finetune (single model.safetensors with `_groot_model.` prefix, " - "config.json type=groot, policy_*processor_step_*.safetensors min/max stats)." - ) - ap.add_argument( - "--stats-out", - type=Path, - default=None, - help="[lerobot ckpt] where to write the bridge's dataset_statistics.json " - "(default: /dataset_statistics.json). state+action min/max are " - "read from the lerobot processor safetensors; the eval bridge consumes it " - "via --stats-json (the un-normalize is a host-side affine, not a ggml concern)." - ) - args = ap.parse_args() +def convert(ckpt: Path, out: Path, *, writer_factory=open_writer, stats_out: Path | None = None) -> Path: + """Write the GGUF for the checkpoint at `ckpt` to `out` and return `out`. - ckpt = args.ckpt.resolve() - out = resolve_out(args, ckpt, ARCH) + `writer_factory(out, arch)` supplies the gguf.GGUFWriter; the quantized-model converter passes one + that rewrites the FoldQuant sites (docs/QUANTIZATION.md) as they are added. + `stats_out` is --stats-out (lerobot checkpoints only).""" + ckpt = ckpt.resolve() + out = out.resolve() cfg_json = read_json(ckpt / "config.json") if not (ckpt / "model.safetensors.index.json").exists() and not (ckpt / "model.safetensors").exists(): raise SystemExit(f"no model.safetensors[.index.json] under {ckpt}") @@ -223,7 +209,7 @@ def main() -> int: f"horizon={AH['action_horizon']} action_dim={AH['action_dim']} max_state={AH['max_state_dim']} N_steps={AH['num_inference_timesteps']} " f"future_tok={AH['num_target_vision_tokens']} embodiments={AH['max_num_embodiments']} metadata.json={len(metadata_json)} chars") - writer = open_writer(out, ARCH) + writer = writer_factory(out, ARCH) kv_u32( writer, KV, @@ -295,12 +281,35 @@ def main() -> int: write_dit_blocks(writer, g, f"{AHK}.model.transformer_blocks", "aex.dit", AH["dit_layers"]) write_gr00t_proj_out(writer, g, AHK, "aex.dit") - rc = finish(writer, out, " - combined GGUF (Eagle-2.5-VL + action head + cfg + metadata.json)") + finish(writer, out, " - combined GGUF (Eagle-2.5-VL + action head + cfg + metadata.json)") if is_lerobot: - stats_out = (args.stats_out or out.parent / "dataset_statistics.json").resolve() + stats_out = (stats_out or out.parent / "dataset_statistics.json").resolve() _write_lerobot_stats(ckpt, stats_out) - return rc + return out + +def main() -> int: + ap = arg_parser( + ARCH, + "GR00T-N1.5-3B checkpoint dir. Two layouts are auto-detected: " + "(a) NVIDIA Isaac snapshot (sharded safetensors, un-prefixed tensor names, " + "config.json model_type=gr00t_n1_5, experiment_cfg/metadata.json); " + "(b) lerobot finetune (single model.safetensors with `_groot_model.` prefix, " + "config.json type=groot, policy_*processor_step_*.safetensors min/max stats)." + ) + ap.add_argument( + "--stats-out", + type=Path, + default=None, + help="[lerobot ckpt] where to write the bridge's dataset_statistics.json " + "(default: /dataset_statistics.json). state+action min/max are " + "read from the lerobot processor safetensors; the eval bridge consumes it " + "via --stats-json (the un-normalize is a host-side affine, not a ggml concern)." + ) + args = ap.parse_args() + ckpt = args.ckpt.resolve() + convert(ckpt, resolve_out(args, ckpt, ARCH), stats_out=args.stats_out) + return 0 if __name__ == "__main__": raise SystemExit(main()) diff --git a/scripts/convert_gr00t_n1_6_to_gguf.py b/scripts/convert_gr00t_n1_6_to_gguf.py index 8e078bd..d1f6e41 100644 --- a/scripts/convert_gr00t_n1_6_to_gguf.py +++ b/scripts/convert_gr00t_n1_6_to_gguf.py @@ -16,6 +16,7 @@ from __future__ import annotations import json +from pathlib import Path import torch import torch.nn.functional as F @@ -107,29 +108,18 @@ def _resize_pos_embd(pos: torch.Tensor, grid: int) -> torch.Tensor: f"({native * native}→{grid * grid}) bilinear+antialias") return p.reshape(pos.shape[-1], grid * grid).transpose(0, 1).contiguous().to(src_dtype) -def main() -> int: - ap = arg_parser(ARCH, "GR00T-N1.6-3B snapshot dir") - ap.add_argument( - "--vision-size", - type=int, - default=None, - help="Override the SigLIP2 vision-tower input resolution (default: native 224). " - "Set 252 to match the reference processor's smart_resize(factor=28) of a " - "256px image (252 = 18×14 patches ⇒ 324 patches ⇒ 81 tokens after " - "pixel_shuffle÷2). When != 224 the `vit.pos_embd` is bilinear-antialias " - "interpolated from the native 16×16 grid to the new grid, exactly mirroring " - "SiglipVisionEmbeddings.resize_positional_embeddings (F.interpolate " - "mode=bilinear, align_corners=False, antialias=True, float32). The runtime " - "is otherwise resolution-agnostic (grid = image_size/patch_size)." - ) - args = ap.parse_args() - if args.vision_size is not None: - if args.vision_size % VIT["patch_size"] != 0: - raise SystemExit(f"--vision-size {args.vision_size} not divisible by patch_size {VIT['patch_size']}") - VIT["image_size"] = int(args.vision_size) +def convert(ckpt: Path, out: Path, *, writer_factory=open_writer, vision_size: int | None = None) -> Path: + """Write the GGUF for the checkpoint at `ckpt` to `out` and return `out`. - ckpt = args.ckpt.resolve() - out = resolve_out(args, ckpt, ARCH) + `writer_factory(out, arch)` supplies the gguf.GGUFWriter; the quantized-model converter passes one + that rewrites the FoldQuant sites (docs/QUANTIZATION.md) as they are added. + `vision_size` is --vision-size (see main).""" + if vision_size is not None: + if vision_size % VIT["patch_size"] != 0: + raise SystemExit(f"--vision-size {vision_size} not divisible by patch_size {VIT['patch_size']}") + VIT["image_size"] = int(vision_size) + ckpt = ckpt.resolve() + out = out.resolve() require(ckpt / "model.safetensors.index.json") cfg_json = read_json(ckpt / "config.json") if str(cfg_json.get("model_type", "")) != "Gr00tN1d6": @@ -193,7 +183,7 @@ def main() -> int: f"embodiments={AH['max_num_embodiments']} relative_action={USE_RELATIVE_ACTION} percentiles={USE_PERCENTILES} clip_outliers={CLIP_OUTLIERS} sincos_state={APPLY_SINCOS_STATE} " f"stats={len(statistics_json)}c proc={len(processor_json)}c emb_id={embodiment_id_json.strip()}") - writer = open_writer(out, ARCH) + writer = writer_factory(out, ARCH) kv_u32( writer, KV, @@ -273,7 +263,28 @@ def main() -> int: write_dit_blocks(writer, g, f"{AHK}.model.transformer_blocks", "aex.dit", AH["dit_layers"]) write_gr00t_proj_out(writer, g, AHK, "aex.dit") - return finish(writer, out, " - combined GGUF (Eagle-3-VL + AlternateVLDiT action head + cfg + sidecars)") + finish(writer, out, " - combined GGUF (Eagle-3-VL + AlternateVLDiT action head + cfg + sidecars)") + return out + +def main() -> int: + ap = arg_parser(ARCH, "GR00T-N1.6-3B snapshot dir") + ap.add_argument( + "--vision-size", + type=int, + default=None, + help="Override the SigLIP2 vision-tower input resolution (default: native 224). " + "Set 252 to match the reference processor's smart_resize(factor=28) of a " + "256px image (252 = 18×14 patches ⇒ 324 patches ⇒ 81 tokens after " + "pixel_shuffle÷2). When != 224 the `vit.pos_embd` is bilinear-antialias " + "interpolated from the native 16×16 grid to the new grid, exactly mirroring " + "SiglipVisionEmbeddings.resize_positional_embeddings (F.interpolate " + "mode=bilinear, align_corners=False, antialias=True, float32). The runtime " + "is otherwise resolution-agnostic (grid = image_size/patch_size)." + ) + args = ap.parse_args() + ckpt = args.ckpt.resolve() + convert(ckpt, resolve_out(args, ckpt, ARCH), vision_size=args.vision_size) + return 0 if __name__ == "__main__": raise SystemExit(main()) diff --git a/scripts/convert_gr00t_n1_7_to_gguf.py b/scripts/convert_gr00t_n1_7_to_gguf.py index 7009f9f..395b2c5 100644 --- a/scripts/convert_gr00t_n1_7_to_gguf.py +++ b/scripts/convert_gr00t_n1_7_to_gguf.py @@ -16,6 +16,7 @@ from __future__ import annotations import json +from pathlib import Path from gguf_blocks import ( write_dit_blocks, @@ -143,12 +144,15 @@ def _uses_relative_actions(ckpt, processor_json: str) -> bool: LM_ROOT = "backbone.model.model.language_model" AHK = "action_head" -def main() -> int: - ap = arg_parser(ARCH, "GR00T-N1.7-3B snapshot dir") - args = ap.parse_args() +def convert(ckpt: Path, out: Path, *, writer_factory=open_writer) -> Path: + """Write the GGUF for the checkpoint at `ckpt` to `out` and return `out`. - ckpt = args.ckpt.resolve() - out = resolve_out(args, ckpt, ARCH) + `writer_factory(out, arch)` supplies the gguf.GGUFWriter; the quantized-model converter passes one + that rewrites the FoldQuant sites (docs/QUANTIZATION.md) as they are added, + so the exported file is this converter's file with INT8 sites, not a + re-implementation of it.""" + ckpt = ckpt.resolve() + out = out.resolve() require(ckpt / "model.safetensors.index.json") cfg_json = read_json(ckpt / "config.json") if str(cfg_json.get("model_type", "")) != "Gr00tN1d7": @@ -233,7 +237,7 @@ def main() -> int: f"embodiments={AH['max_num_embodiments']} relative={USE_RELATIVE_ACTION} percentiles={USE_PERCENTILES} clip={CLIP_OUTLIERS} sincos={APPLY_SINCOS_STATE} " f"img: shortest_edge={SHORTEST_EDGE} crop_frac={CROP_FRACTION} crop_size={ICS} target_size={ITS} stats={len(statistics_json)}c proc={len(processor_json)}c emb_id={embodiment_id_json.strip()}") - writer = open_writer(out, ARCH) + writer = writer_factory(out, ARCH) kv_u32( writer, KV, @@ -330,12 +334,20 @@ def main() -> int: write_dit_blocks(writer, g, f"{AHK}.model.transformer_blocks", "aex.dit", AH["dit_layers"]) write_gr00t_proj_out(writer, g, AHK, "aex.dit") - return finish( + finish( writer, out, " - combined GGUF (Qwen3-VL backbone + deepstack + vl_self_attention " "+ AlternateVLDiT action head + cfg + sidecars)" ) + return out + +def main() -> int: + ap = arg_parser(ARCH, "GR00T-N1.7-3B snapshot dir") + args = ap.parse_args() + ckpt = args.ckpt.resolve() + convert(ckpt, resolve_out(args, ckpt, ARCH)) + return 0 if __name__ == "__main__": raise SystemExit(main()) diff --git a/scripts/convert_pi05_to_gguf.py b/scripts/convert_pi05_to_gguf.py index 565ae26..bf387e0 100644 --- a/scripts/convert_pi05_to_gguf.py +++ b/scripts/convert_pi05_to_gguf.py @@ -122,6 +122,9 @@ def grab(feat, dim): out[q] = np.asarray(s[q], dtype=np.float32).reshape(-1)[:dim] return out + if "norm_stats" in d: + # OpenPI assets///norm_stats.json: {"norm_stats": {"state": .., "actions": ..}} + d = {"observation.state": d["norm_stats"]["state"], "action": d["norm_stats"]["actions"]} st = grab("observation.state", state_dim) ac = grab("action", action_dim) return { @@ -150,30 +153,32 @@ def grab(feat, dim): "--dataset-stats or --dataset-repo lerobot/libero. " "Refusing to bake identity stats (would make the policy miss; skill footgun #1).") -def main() -> int: - ap = arg_parser(ARCH, "lerobot π0.5 checkpoint dir (model.safetensors + config.json + policy_*processor.json)") - ap.add_argument( - "--dataset-stats", - type=Path, - default=None, - help="Path to a LIBERO dataset meta/stats.json for QUANTILES norm stats" - ) - ap.add_argument( - "--dataset-repo", - type=str, - default=None, - help="HF dataset repo to fetch meta/stats.json from (e.g. lerobot/libero)" - ) - args = ap.parse_args() - - ckpt = args.ckpt.resolve() - out = resolve_out(args, ckpt, ARCH) +def convert(ckpt: Path, out: Path, *, writer_factory=open_writer, dataset_stats: Optional[Path] = None, + dataset_repo: Optional[str] = None, config: Optional[dict] = None) -> Path: + """Write the GGUF for the checkpoint at `ckpt` to `out` and return `out`. + + `writer_factory(out, arch)` supplies the gguf.GGUFWriter; the quantized-model converter passes one + that rewrites the FoldQuant sites (docs/QUANTIZATION.md) as they are added. + `dataset_stats` / `dataset_repo` are --dataset-stats / --dataset-repo (a + LeRobot meta/stats.json or an OpenPI norm_stats.json). `config` supplies the + lerobot policy fields (chunk_size, num_inference_steps, n_action_steps, + max_state_dim, max_action_dim, min_period, max_period, tokenizer_max_length, + input_features, output_features) for an OpenPI-converted checkpoint whose + config.json carries none of them (--config-json).""" + ckpt = ckpt.resolve() + out = out.resolve() sf_path = ckpt / "model.safetensors" require(sf_path) - cfg_json = read_json(ckpt / "config.json") + # An OpenPI checkpoint may carry no config.json at all; `config` then supplies everything. + cfg_json = read_json(ckpt / "config.json") if (ckpt / "config.json").exists() or config is None else {} if cfg_json.get("type") != ARCH: - raise SystemExit(f"config.json type is {cfg_json.get('type')!r}, expected 'pi05'") + if config is None: + raise SystemExit(f"config.json type is {cfg_json.get('type')!r}, expected 'pi05'") + # An OpenPI-converted checkpoint (config.json carries only the variants): + # the policy fields come from the caller (the quantized model's export metadata, or + # --config-json), in the lerobot field names. + cfg_json = {**cfg_json, **config} norm_map = cfg_json.get("normalization_mapping") or {} norm_mode = norm_map.get("ACTION", "QUANTILES") if norm_mode != "QUANTILES" or norm_map.get("STATE", "QUANTILES") != norm_mode: @@ -249,15 +254,15 @@ def main() -> int: print("loading dataset normalizer stats...") stats = _load_dataset_stats( - args.dataset_stats, - args.dataset_repo, + dataset_stats, + dataset_repo, cfg["real_state_dim"], cfg["real_action_dim"] ) print(f" state_q01[:3]={stats['state_q01'][:3]} state_q99[:3]={stats['state_q99'][:3]}") print(f" action_q01[:3]={stats['action_q01'][:3]} action_q99[:3]={stats['action_q99'][:3]} (QUANTILES)") - writer = open_writer(out, ARCH) + writer = writer_factory(out, ARCH) write_pi_kv(writer, KV, cfg, adarms=True) add(writer, "token_embd.weight", sf.get_tensor(head)) @@ -276,10 +281,38 @@ def main() -> int: for name, vec in stats.items(): add_array(writer, name, vec) - rc = finish(writer, out) + finish(writer, out) print("note: self-contained GGUF — SigLIP vision tower + PaliGemma projector are baked in; " "no separate mmproj is needed.") - return rc + return out + +def main() -> int: + ap = arg_parser(ARCH, "lerobot π0.5 checkpoint dir (model.safetensors + config.json + policy_*processor.json)") + ap.add_argument( + "--dataset-stats", + type=Path, + default=None, + help="Path to a LIBERO dataset meta/stats.json for QUANTILES norm stats" + ) + ap.add_argument( + "--dataset-repo", + type=str, + default=None, + help="HF dataset repo to fetch meta/stats.json from (e.g. lerobot/libero)" + ) + ap.add_argument( + "--config-json", + type=Path, + default=None, + help="[OpenPI-converted checkpoint] JSON with the lerobot policy fields its config.json lacks " + "(chunk_size, num_inference_steps, n_action_steps, max_state_dim, max_action_dim, " + "min_period, max_period, tokenizer_max_length, input_features, output_features)" + ) + args = ap.parse_args() + ckpt = args.ckpt.resolve() + convert(ckpt, resolve_out(args, ckpt, ARCH), dataset_stats=args.dataset_stats, dataset_repo=args.dataset_repo, + config=read_json(args.config_json) if args.config_json else None) + return 0 if __name__ == "__main__": raise SystemExit(main()) diff --git a/scripts/convert_quantized_model_to_gguf.py b/scripts/convert_quantized_model_to_gguf.py new file mode 100644 index 0000000..986f342 --- /dev/null +++ b/scripts/convert_quantized_model_to_gguf.py @@ -0,0 +1,1152 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Convert a FoldQuant quantized model to a vla.cpp FoldQuant GGUF. + +A quantized model is a quantized policy saved as data instead of as an engine: +the base checkpoint plus the calibration result and the integer weight codes of +every quantized site, exactly as the TensorRT plugins would carry them. +FoldQuantVLA writes one in either of two formats, and this converter reads both: + +* FoldQuantVLA's quantized checkpoint, `foldquant.quantized_checkpoint` + (format "foldquant-quantized-checkpoint"): the base checkpoint with every + quantized projection's `.weight` replaced by `.qweight` (int8, or int4 + nibble-packed) + `.weight_scale`, the SmoothQuant vectors as + `foldquant..sq.` tensors, and `foldquant_quant.json` naming the + checkpoint keys of every site. The family converter reads the directory + through a base view in which each missing `.weight` is qweight x scale; the + sites are then written as codes, and the one quantized projection vla.cpp + keeps in float, the DiT adaLN of a W4A4 arm, from that dequantized weight. +* FoldQuantVLA's earlier fake-quant state, `foldquant.fakequant` (format + "foldquant-quant-state"): the base checkpoint files next to + `foldquant_quant.json` and `quant_state.safetensors`, sites by name, INT8 DiT + packs by call order. + + For both, the LLM's SmoothQuant fold into the norm gains is recomputed here + the way FoldQuant's emitter computes it. + +Nothing is recalibrated or re-rounded: the recorded codes go into the GGUF as +they are (FoldQuant's INT4 layout is vla.cpp's: row-major, low nibble = even +column), each action site ships its SmoothQuant vector as `.ascale`, and the +rest of the file is the family's own converter output. + + python scripts/convert_quantized_model_to_gguf.py --quantized-model foldquant_model --out model.gguf + python scripts/convert_quantized_model_to_gguf.py --quantized-model foldquant_model --out model.gguf --check-onnx exports/arm/onnx + +`--check-onnx` byte-compares every site and folded gain +against the plugin ONNX graphs the TensorRT engines were built from and fails +on any difference. + +Supported families: GR00T N1.5 / N1.6 / N1.7 (Qwen3 LLM + DiT) and pi0.5 +(PaliGemma prefix + Gemma expert). Supported schemes are the ones vla.cpp has +kernels for: `w8a8_sr` / `w4a4_srg` LLM (with `site_bits`) and `w8a8_sh` / +`w4a4_sh` / `w4a4_shg` action arms. Action arms with a dense learned rotation +(`*_sr`) are refused. See docs/QUANTIZATION.md. +""" + +from __future__ import annotations + +import argparse +import contextlib +import importlib +import json +import re +import sys +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional, Tuple + +import numpy as np +import torch +from safetensors import safe_open + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from gguf_quant_writer import QuantKV, Site, quantizing_writer_factory # noqa: E402 + +# FoldQuantVLA: both formats keep their manifest under this name; `format` tells them apart. +FQ_MANIFEST = "foldquant_quant.json" +FQ_TENSORS = "quant_state.safetensors" +FQ_FORMAT = "foldquant-quant-state" # base checkpoint + a separate quant state +FQC_FORMAT = "foldquant-quantized-checkpoint" # the base checkpoint with .qweight/.weight_scale in place +FQC_MAX_VERSION = 1 + +LLM_SCHEMES = {"w8a8_sr", "w4a4_srg"} +ACTION_SCHEMES = {"w8a8_sh", "w4a4_sh", "w4a4_shg"} + +LLM_SITES = ("qkv", "o", "gateup", "down") +LLM_NODES = {"qkv": "L{i}_rmsnorm_qkv", "o": "L{i}_o_proj_res", "gateup": "L{i}_rmsnorm_gateup", "down": "L{i}_down_proj_res"} +EXPERT_SITES = ("qkv", "o", "gu", "dn") + + +@dataclass +class Rec: + """One recorded site: codes int8 (N, K) or uint8 (N, K/2), per-row scale (N,), SmoothQuant vector (K,) or None.""" + + codes: np.ndarray + scale: np.ndarray + ascale: Optional[np.ndarray] = None + + @property + def bits(self) -> int: + return 4 if self.codes.dtype == np.uint8 else 8 + + @property + def n_out(self) -> int: + return int(self.codes.shape[0]) + + def split(self, rows: List[int]) -> List[Tuple[np.ndarray, np.ndarray]]: + if sum(rows) != self.n_out: + raise SystemExit(f"recorded codes have {self.n_out} rows, the projections {rows}") + out, r0 = [], 0 + for n in rows: + out.append((np.ascontiguousarray(self.codes[r0:r0 + n]), np.ascontiguousarray(self.scale[r0:r0 + n]))) + r0 += n + return out + + +def _as_codes(buf: bytes, n: int, k: int, bits: int) -> np.ndarray: + if bits == 4: + if len(buf) != n * k // 2: + raise SystemExit(f"INT4 weight blob has {len(buf)} bytes, expected {n}x{k}/2") + return np.frombuffer(buf, np.uint8).reshape(n, k // 2).copy() + if len(buf) != n * k: + raise SystemExit(f"INT8 weight blob has {len(buf)} bytes, expected {n}x{k}") + return np.frombuffer(buf, np.int8).reshape(n, k).copy() + + +def _fold_gain(gain: torch.Tensor, s: np.ndarray, *, gemma: bool) -> torch.Tensor: + """FoldQuant's SmoothQuant fold of a norm gain, as its plugin emitter writes it: (w [+ 1]) / s in fp32, + rounded once to the gain's own dtype (byte-checked against the TensorRT graphs' `gamma`).""" + dt = gain.dtype + g32 = gain if dt in (torch.float32, torch.float64) else gain.float() + if gemma: + g32 = g32 + 1.0 # Gemma applies (1 + w); the 1 is added in fp32, not in bf16 + return (g32 / torch.from_numpy(np.asarray(s, np.float32))).to(dt) + + +def _gamma_from_bytes(buf: bytes, k: int) -> torch.Tensor: + if len(buf) == 2 * k: + return torch.frombuffer(bytearray(buf), dtype=torch.bfloat16).clone() + if len(buf) == 4 * k: + return torch.frombuffer(bytearray(buf), dtype=torch.float32).clone() + raise SystemExit(f"gamma blob has {len(buf)} bytes for K={k}: neither bf16 nor f32") + + +# --------------------------------------------------------------------------- +# base checkpoint +# --------------------------------------------------------------------------- + + +def _unpack_int4(packed: torch.Tensor) -> torch.Tensor: + """uint8 (N, K/2), low nibble = even column, two's complement -> int8 (N, K).""" + p = packed.to(torch.int16) + both = torch.stack([p & 0xF, (p >> 4) & 0xF], dim=-1).reshape(p.shape[0], p.shape[1] * 2) + return torch.where(both >= 8, both - 16, both).to(torch.int8) + + +def _is_adaln(site: str) -> bool: + return site.endswith("_adaln") + + +class Checkpoint: + """The base checkpoint's safetensors (one file or a sharded index). + + A FoldQuantVLA quantized checkpoint (FQC_FORMAT) is the base checkpoint with every quantized + projection's `

.weight` replaced by `

.qweight` + `

.weight_scale`, plus `foldquant.*` site + tensors. Opened on one, this class presents the base view the family converters expect: the + `.qweight` / `.weight_scale` / `foldquant.*` tensors are hidden and each `

.weight` is rebuilt as + qweight x weight_scale (bf16). A site's rebuilt weight is only read for its shape (the quantizing + writer replaces it with the codes); the one quantized projection vla.cpp keeps in float, the DiT + adaLN of a W4A4 arm, is written from it, with the scale rounded to bf16 first as the TensorRT + AdaLNModInt4 plugin bakes it. + """ + + def __init__(self, root: Path) -> None: + self.root = root.resolve() + index = self.root / "model.safetensors.index.json" + single = self.root / "model.safetensors" + self._files: Dict[str, Path] = {} + if index.is_file(): + for key, fname in json.loads(index.read_text())["weight_map"].items(): + self._files[key] = self.root / fname + elif single.is_file(): + with safe_open(str(single), framework="pt") as f: + for key in f.keys(): + self._files[key] = single + else: + raise SystemExit(f"{self.root}: no model.safetensors or model.safetensors.index.json; pass --ckpt ") + self._open: Dict[Path, Any] = {} + # "

.weight" -> (bits, is_adaln) for a quantized checkpoint's replaced projections + self.quantized: Dict[str, Tuple[int, bool]] = {} + manifest = self.root / FQ_MANIFEST + if manifest.is_file(): + m = json.loads(manifest.read_text()) + if m.get("format") == FQC_FORMAT: + for module in (m.get("modules") or {}).values(): + for key, info in (module.get("sites") or {}).items(): + for group in info["params"]: + for ck in group: + self.quantized[ck] = (int(info["bits"]), _is_adaln(key)) + hidden = set() + for ck in self.quantized: + head = ck[: -len(".weight")] + for suffix in (".qweight", ".weight_scale"): + if head + suffix not in self._files: + raise SystemExit(f"{self.root}: quantized projection {head} has no {head}{suffix}") + hidden.add(head + suffix) + hidden |= {k for k in self._files if k.startswith("foldquant.")} + self._hidden = hidden + + @property + def is_quantized(self) -> bool: + return bool(self.quantized) + + @property + def weight_files(self) -> List[Path]: + return sorted(set(self._files.values())) + + def _f(self, key: str) -> Any: + if key not in self._files: + raise SystemExit(f"{self.root}: no tensor {key!r} in the base checkpoint") + p = self._files[key] + if p not in self._open: + self._open[p] = safe_open(str(p), framework="pt") + return self._open[p] + + def raw(self, key: str) -> torch.Tensor: + """A tensor as stored, hidden ones included.""" + return self._f(key).get_tensor(key) + + def keys(self) -> List[str]: + """The base view's tensor names.""" + return sorted((set(self._files) - self._hidden) | set(self.quantized)) + + def shape(self, key: str) -> Tuple[int, ...]: + if key in self.quantized: + head, (bits, _) = key[: -len(".weight")], self.quantized[key] + n, cols = self._f(head + ".qweight").get_slice(head + ".qweight").get_shape() + return (int(n), int(cols) * (2 if bits == 4 else 1)) + if key in self._hidden: + raise SystemExit(f"{self.root}: {key!r} is part of a quantized projection, not a base tensor") + return tuple(self._f(key).get_slice(key).get_shape()) + + def tensor(self, key: str) -> torch.Tensor: + if key in self.quantized: + head, (bits, adaln) = key[: -len(".weight")], self.quantized[key] + codes = self.raw(head + ".qweight") + codes = _unpack_int4(codes) if bits == 4 else codes.to(torch.int8) + scale = self.raw(head + ".weight_scale").float() + if adaln: + scale = scale.to(torch.bfloat16).float() + return (codes.float() * scale[:, None]).to(torch.bfloat16) + if key in self._hidden: + raise SystemExit(f"{self.root}: {key!r} is part of a quantized projection, not a base tensor") + return self._f(key).get_tensor(key) + + # -- what the family converters read ------------------------------------------------ + + def mapping(self, keep: Optional[Tuple[str, ...]] = None) -> "_TensorMap": + return _TensorMap(self, keep) + + def handle(self) -> "_Handle": + return _Handle(self) + + +class _TensorMap(Mapping): + """`gguf_common.load_safetensors(ckpt)` over the base view, read lazily.""" + + def __init__(self, ckpt: Checkpoint, keep: Optional[Tuple[str, ...]]) -> None: + self._ckpt = ckpt + self._keys = [k for k in ckpt.keys() if keep is None or k.startswith(keep)] + self._set = set(self._keys) + + def __getitem__(self, key: str) -> torch.Tensor: + if key not in self._set: + raise KeyError(key) + return self._ckpt.tensor(key) + + def __iter__(self) -> Iterator[str]: + return iter(self._keys) + + def __len__(self) -> int: + return len(self._keys) + + +class _Slice: + def __init__(self, shape: Tuple[int, ...]) -> None: + self._shape = list(shape) + + def get_shape(self) -> List[int]: + return list(self._shape) + + +class _Handle: + """`safe_open(/model.safetensors)` over the base view (keys, get_tensor, get_slice().get_shape()).""" + + def __init__(self, ckpt: Checkpoint) -> None: + self._ckpt = ckpt + + def keys(self) -> List[str]: + return self._ckpt.keys() + + def get_tensor(self, key: str) -> torch.Tensor: + return self._ckpt.tensor(key) + + def get_slice(self, key: str) -> _Slice: + return _Slice(self._ckpt.shape(key)) + + def metadata(self) -> Dict[str, str]: + return {"format": "pt"} + + def __enter__(self) -> "_Handle": + return self + + def __exit__(self, *exc: Any) -> None: + return None + + +@contextlib.contextmanager +def base_view(conv: Any, ckpt: Checkpoint) -> Iterator[None]: + """Point a family converter's checkpoint reads at `ckpt`'s base view while it converts. + + The GR00T converters read through `load_safetensors` and the pi0.5 converter through `safe_open`, + both imported into the converter module by name; reads of any other path pass through. + """ + if not ckpt.is_quantized: + yield + return + files = {p.resolve() for p in ckpt.weight_files} + saved: Dict[str, Any] = {} + if hasattr(conv, "load_safetensors"): + real_load = saved["load_safetensors"] = conv.load_safetensors + + def load_safetensors(path: Any, keep: Optional[Tuple[str, ...]] = None) -> Any: + if Path(path).resolve() == ckpt.root: + return ckpt.mapping(keep) + return real_load(path, keep) + + conv.load_safetensors = load_safetensors + if hasattr(conv, "safe_open"): + real_open = saved["safe_open"] = conv.safe_open + + def open_(path: Any, framework: str = "pt", device: str = "cpu") -> Any: + if Path(path).resolve() in files: + if len(files) != 1: + raise SystemExit(f"{ckpt.root}: a sharded quantized checkpoint read through safe_open() by the converter") + return ckpt.handle() + return real_open(path, framework=framework, device=device) + + conv.safe_open = open_ + try: + yield + finally: + for name, value in saved.items(): + setattr(conv, name, value) + + +# --------------------------------------------------------------------------- +# sources +# --------------------------------------------------------------------------- + + +class Source: + """What a quantized-model producer recorded: schemes, sites, folded gains.""" + + producer: str + family: str + root: Path + base_checkpoint: Optional[Path] + + def scheme(self, module: str) -> Tuple[str, Dict[str, Any]]: + raise NotImplementedError + + def llm_layers(self) -> int: + raise NotImplementedError + + def llm_site(self, i: int, key: str) -> Rec: + raise NotImplementedError + + def llm_gamma(self, i: int, key: str, ckpt_gain: torch.Tensor, *, gemma: bool) -> torch.Tensor: + """The SmoothQuant-folded norm gain feeding site `key` (qkv -> attn norm, gateup -> ffn norm).""" + raise NotImplementedError + + def llm_rot_block(self) -> int: + raise NotImplementedError + + def dit_blocks(self) -> int: + raise NotImplementedError + + def dit_block(self, i: int) -> Dict[str, Rec]: + """Sites of DiT block i: {qkv, o, ffn0, ffn2} (self-attention) or {q, kv, o, ffn0, ffn2} (cross).""" + raise NotImplementedError + + def dit_rot_block(self) -> int: + return 64 + + def expert_layers(self) -> int: + raise NotImplementedError + + def expert_site(self, i: int, key: str) -> Rec: + raise NotImplementedError + + def provenance(self) -> str: + raise NotImplementedError + + +class FoldQuantState(Source): + """FoldQuantVLA's quantized model: named site records + SmoothQuant vectors.""" + + producer = "foldquant" + + def __init__(self, root: Path) -> None: + self.root = root.resolve() + self.manifest = json.loads((self.root / FQ_MANIFEST).read_text()) + if self.manifest.get("format") != FQ_FORMAT: + raise SystemExit(f"{self.root / FQ_MANIFEST}: format {self.manifest.get('format')!r}, expected {FQ_FORMAT!r}") + self.family = str(self.manifest["family"]) + self.base_checkpoint = self.root # the state is saved next to (or bundled with) the base files + self._st = safe_open(str(self.root / FQ_TENSORS), framework="np") + self.keys = set(self._st.keys()) + + def scheme(self, module: str) -> Tuple[str, Dict[str, Any]]: + mod = self.manifest["modules"].get(module) + if mod is None: + raise SystemExit(f"{self.root / FQ_MANIFEST}: no module {module!r} (has {sorted(self.manifest['modules'])})") + return str(mod["scheme"]), dict(mod.get("config") or {}) + + def tensor(self, key: str) -> np.ndarray: + if key not in self.keys: + raise SystemExit(f"{self.root / FQ_TENSORS}: no tensor {key!r}") + return self._st.get_tensor(key) + + def _pack(self, module: str, group: str, index: int = 0) -> Optional[Tuple[np.ndarray, np.ndarray]]: + ck = f"{module}/g/{group}/{index}/codes" + if ck not in self.keys: + return None + return self.tensor(ck), self.tensor(f"{module}/g/{group}/{index}/scale").astype(np.float32) + + def site(self, module: str, key: str) -> Tuple[np.ndarray, np.ndarray]: + """(codes, scale) of a site recorded under its own name (GPTQ, or site-scoped round-to-nearest).""" + for group in (f"{module}.{key}", f"{module}.rtn.{key}"): + rec = self._pack(module, group) + if rec is not None: + if f"{module}/g/{group}/1/codes" in self.keys: + raise SystemExit(f"{group}: more than one recorded pack; not a single-site record") + return rec + raise SystemExit(f"{self.root / FQ_TENSORS}: site {module}.{key} has no recorded codes " + f"(recorded by call order only, or not quantized)") + + # LLM + def llm_layers(self) -> int: + n = 0 + while f"llm/t/sq/L{n}_qkv" in self.keys: + n += 1 + return n + + def llm_site(self, i: int, key: str) -> Rec: + codes, scale = self.site("llm", f"L{i}_{key}") + return Rec(codes, scale) + + @property + def model_dtype(self) -> torch.dtype: + """The dtype FoldQuant loaded the model in (its export precision; bf16 unless recorded otherwise).""" + prec = str(self.export_metadata().get("precision", "bf16")).lower() + return {"bf16": torch.bfloat16, "bfloat16": torch.bfloat16, "fp16": torch.float16, "float16": torch.float16, + "fp32": torch.float32, "float32": torch.float32}.get(prec, torch.bfloat16) + + def llm_gamma(self, i: int, key: str, ckpt_gain: torch.Tensor, *, gemma: bool) -> torch.Tensor: + # The emitter folds the gain the model holds, i.e. the checkpoint value cast to the export dtype + # (some GR00T checkpoints store fp32 norm gains; the exported model is bf16). + return _fold_gain(ckpt_gain.to(self.model_dtype), self.tensor(f"llm/t/sq/L{i}_{key}"), gemma=gemma) + + def llm_rot_block(self) -> int: + return int(self.scheme("llm")[1].get("rot_bs", 64)) + + # DiT (GR00T) + def dit_blocks(self) -> int: + n = 0 + while f"dit/t/sq/block{n}_o" in self.keys: + n += 1 + return n + + def _dit_is_self(self, i: int) -> bool: + return f"dit/t/sq/block{i}_qkv" in self.keys + + # The INT8 DiT emitter packs round-to-nearest codes in this order per block (FoldQuant records + # them as `dit.rtn/{n}` by call order); the GPTQ INT4 emitter records each site by name instead. + _RTN_ORDER_CROSS = ("q", "o", "kv", "ffn0", "ffn2") + _RTN_ORDER_SELF = ("qkv", "o", "ffn0", "ffn2") + + def dit_block(self, i: int) -> Dict[str, Rec]: + def sq(key: str) -> np.ndarray: + return self.tensor(f"dit/t/sq/{key}").astype(np.float32) + + is_self = self._dit_is_self(i) + names = self._RTN_ORDER_SELF if is_self else self._RTN_ORDER_CROSS + vec = {k: sq("encoder" if k == "kv" else f"block{i}_{k}") for k in names} + if self._pack("dit", f"dit.block{i}_o") is not None: + # GPTQ: every site by name; the cross-attention KV packs share the encoder input and + # are recorded under `dit.encoder` in block order (the j-th record = the j-th cross block). + out: Dict[str, Rec] = {} + for k in names: + if k == "kv": + j = sum(1 for b in range(i) if not self._dit_is_self(b)) + kv = self._pack("dit", "dit.encoder", j) + if kv is None: + raise SystemExit(f"{self.root / FQ_TENSORS}: no dit.encoder pack #{j} for cross block {i}") + out[k] = Rec(kv[0], kv[1], vec[k]) + else: + codes, scale = self.site("dit", f"block{i}_{k}") + out[k] = Rec(codes, scale, vec[k]) + return out + # Round-to-nearest: walk `dit.rtn/{n}` in emitter order, checking every shape against the vectors. + start = sum(len(self._RTN_ORDER_SELF if self._dit_is_self(b) else self._RTN_ORDER_CROSS) for b in range(i)) + d = int(vec["o"].shape[0]) + want_n = {"q": d, "qkv": 3 * d, "kv": 2 * d, "o": int(vec["qkv" if is_self else "q"].shape[0]), + "ffn0": int(vec["ffn2"].shape[0]), "ffn2": int(vec["ffn0"].shape[0])} + out = {} + for n, k in enumerate(names): + pack = self._pack("dit", "dit.rtn", start + n) + if pack is None: + raise SystemExit(f"{self.root / FQ_TENSORS}: block {i} site {k}: no dit.rtn pack #{start + n} " + "(the state records neither named DiT sites nor the expected call-order packs)") + rec = Rec(pack[0], pack[1], vec[k]) + k_in = int(pack[0].shape[1]) * (2 if rec.bits == 4 else 1) + if rec.n_out != want_n[k] or k_in != int(vec[k].shape[0]): + raise SystemExit(f"{self.root / FQ_TENSORS}: dit.rtn pack #{start + n} is {rec.n_out}x{k_in}, block {i} " + f"site {k} expects {want_n[k]}x{int(vec[k].shape[0])}; the emitter order changed") + out[k] = rec + return out + + # Gemma expert (pi0.5) + def expert_layers(self) -> int: + n = 0 + while f"expert/t/sq/G{n}_qkv" in self.keys: + n += 1 + return n + + def expert_site(self, i: int, key: str) -> Rec: + codes, scale = self.site("expert", f"G{i}_{key}") + return Rec(codes, scale, self.tensor(f"expert/t/sq/G{i}_{key}").astype(np.float32)) + + def provenance(self) -> str: + base = self.manifest.get("base") or {} + return f"foldquant quantized model {self.root.name}; base={base.get('model_id')}; digest={str(base.get('digest', ''))[:16]}" + + def export_metadata(self) -> Dict[str, Any]: + return dict((self.manifest.get("extra_files") or {}).get("export_metadata.json") or {}) + + +_SITE_RE = { + "llm": re.compile(r"^L(\d+)_(qkv|o|gateup|down)$"), + "expert": re.compile(r"^G(\d+)_(qkv|o|gu|dn)$"), + "dit": re.compile(r"^(?:block(\d+)_(qkv|q|kv|o|ffn0|ffn2|adaln)|encoder)$"), +} + + +class FoldQuantCheckpoint(FoldQuantState): + """FoldQuantVLA's quantized checkpoint (FQC_FORMAT, `foldquant.quantized_checkpoint`). + + One directory that loads like the base checkpoint: each quantized projection's `.weight` is + replaced by `.qweight` (int8 (N, K), or int4 nibble-packed (N, K/2), low nibble = even column, + which is vla.cpp's layout) and `.weight_scale` (fp32 (N,)); the manifest's `sites` name, per + recorded pack, the checkpoint keys it covers in row order; the SmoothQuant vectors are + `foldquant..sq.` tensors. The base checkpoint is the directory itself. + """ + + producer = "foldquant" + + def __init__(self, root: Path) -> None: # noqa: D107 (FoldQuantState's reader does not apply) + self.root = root.resolve() + self.manifest = json.loads((self.root / FQ_MANIFEST).read_text()) + if self.manifest.get("format") != FQC_FORMAT: + raise SystemExit(f"{self.root / FQ_MANIFEST}: format {self.manifest.get('format')!r}, expected {FQC_FORMAT!r}") + version = int(self.manifest.get("format_version", 0)) + if version > FQC_MAX_VERSION: + raise SystemExit(f"{self.root}: quantized checkpoint format_version {version} is newer than this converter " + f"reads ({FQC_MAX_VERSION}); update vla.cpp") + self.family = str(self.manifest["family"]) + self.base_checkpoint = self.root + self.ckpt = Checkpoint(self.root) + # module -> site name -> (bits, [[checkpoint keys], ...]) + self.sites: Dict[str, Dict[str, Tuple[int, List[List[str]]]]] = {} + for module, meta in (self.manifest.get("modules") or {}).items(): + table: Dict[str, Tuple[int, List[List[str]]]] = {} + for key, info in (meta.get("sites") or {}).items(): + name = key[len(module) + 1:] if key.startswith(module + ".") else key + name = name[len("rtn."):] if name.startswith("rtn.") else name + pattern = _SITE_RE.get(module) + if pattern is None or not pattern.match(name): + raise SystemExit(f"{self.root}: module {module!r} records site {key!r}, which vla.cpp has no mapping for") + if name in table: + raise SystemExit(f"{self.root}: site {module}.{name} is recorded twice") + table[name] = (int(info["bits"]), [list(g) for g in info["params"]]) + self.sites[module] = table + self.consumed: set = set() # checkpoint keys that went into a GGUF site + self.keys = set() # FoldQuantState's key set is not used here + + def tensor(self, key: str) -> np.ndarray: # "llm/t/sq/L0_qkv" -> foldquant.llm.sq.L0_qkv + module, _, rest = key.partition("/t/") + return self.ckpt.raw(f"foldquant.{module}.{rest.replace('/', '.')}").float().numpy() + + def _has_sq(self, module: str, site: str) -> bool: + return f"foldquant.{module}.sq.{site}" in self.ckpt._files + + def _rec(self, module: str, site: str, group: int = 0) -> Tuple[np.ndarray, np.ndarray]: + entry = self.sites.get(module, {}).get(site) + if entry is None: + raise SystemExit(f"{self.root}: no quantized site {module}.{site}") + bits, groups = entry + if group >= len(groups): + raise SystemExit(f"{self.root}: site {module}.{site} has {len(groups)} group(s), asked for #{group}") + codes, scales = [], [] + for ck in groups[group]: + head = ck[: -len(".weight")] + c = self.ckpt.raw(head + ".qweight") + want = torch.uint8 if bits == 4 else torch.int8 + if c.dtype != want: + raise SystemExit(f"{self.root}: {head}.qweight is {c.dtype}, a {bits}-bit site stores {want}") + codes.append(c.numpy()) + scales.append(self.ckpt.raw(head + ".weight_scale").float().numpy()) + self.consumed.add(ck) + return np.ascontiguousarray(np.concatenate(codes)), np.ascontiguousarray(np.concatenate(scales).astype(np.float32)) + + def site(self, module: str, key: str) -> Tuple[np.ndarray, np.ndarray]: + return self._rec(module, key) + + def scheme(self, module: str) -> Tuple[str, Dict[str, Any]]: + mod = (self.manifest.get("modules") or {}).get(module) + if mod is None: + raise SystemExit(f"{self.root / FQ_MANIFEST}: no module {module!r} (has {sorted(self.manifest.get('modules') or {})})") + return str(mod["scheme"]), dict(mod.get("config") or {}) + + # LLM + def llm_layers(self) -> int: + return sum(1 for s in self.sites.get("llm", {}) if s.endswith("_qkv")) + + # DiT + def dit_blocks(self) -> int: + return sum(1 for s in self.sites.get("dit", {}) if s.startswith("block") and s.endswith("_o")) + + def _dit_is_self(self, i: int) -> bool: + return f"block{i}_qkv" in self.sites.get("dit", {}) + + def dit_block(self, i: int) -> Dict[str, Rec]: + def rec(site: str, sq: str, group: int = 0) -> Rec: + codes, scale = self._rec("dit", site, group) + return Rec(codes, scale, self.tensor(f"dit/t/sq/{sq}").astype(np.float32)) + + out = {k: rec(f"block{i}_{k}", f"block{i}_{k}") for k in ("o", "ffn0", "ffn2")} + if self._dit_is_self(i): + out["qkv"] = rec(f"block{i}_qkv", f"block{i}_qkv") + return out + out["q"] = rec(f"block{i}_q", f"block{i}_q") + # Cross-attention K/V read the encoder, so they share its SmoothQuant vector. Round-to-nearest + # records them per block (`block{i}_kv`); GPTQ under `encoder`, one group per cross block in order. + if f"block{i}_kv" in self.sites.get("dit", {}): + out["kv"] = rec(f"block{i}_kv", "encoder") + else: + j = sum(1 for b in range(i) if not self._dit_is_self(b)) + out["kv"] = rec("encoder", "encoder", j) + return out + + # Gemma expert + def expert_layers(self) -> int: + return sum(1 for s in self.sites.get("expert", {}) if s.endswith("_qkv")) + + def float_projections(self) -> List[str]: + """Quantized projections vla.cpp runs in float, from their dequantized weight: the DiT adaLN.""" + return sorted(ck for site, (_, groups) in self.sites.get("dit", {}).items() if _is_adaln(site) for g in groups for ck in g) + + def check_accounting(self) -> None: + """Every quantized projection is a GGUF site or an allowed float projection; nothing is dropped.""" + left = set(self.ckpt.quantized) - self.consumed - set(self.float_projections()) + if left: + raise SystemExit(f"{self.root}: {len(left)} quantized projection(s) are neither a vla.cpp site nor an allowed " + f"float projection: {sorted(left)[:4]}") + + def provenance(self) -> str: + base = self.manifest.get("base") or {} + adaln = " adaLN dequantized from INT4;" if self.float_projections() else "" + return (f"foldquant quantized checkpoint {self.root.name} (format {FQC_FORMAT} v{self.manifest.get('format_version')});" + f"{adaln} base={base.get('model_id')}; digest={str(base.get('digest', ''))[:16]}") + + +class PluginGraphs(Source): + """Sites read from FoldQuant plugin ONNX graphs (the nodes the TensorRT engines are built from). + + `files` maps a module (`llm`, `dit`, `expert`) to its graph. Used by `--check-onnx`. + """ + + producer = "onnx" + + def __init__(self, files: Dict[str, Path], family: str = "", + schemes: Optional[Dict[str, Tuple[str, Dict[str, Any]]]] = None) -> None: + self.files = {m: Path(p) for m, p in files.items()} + self.family = family + self.root = next(iter(self.files.values())).parent if self.files else Path(".") + self.base_checkpoint = None + self._schemes = schemes or {} + self._nodes: Dict[str, Dict[str, Any]] = {} + + @staticmethod + def find(directory: Path, modules: Tuple[str, ...] = ("llm", "dit", "expert")) -> Dict[str, Path]: + """`{module: graph}` for the `*.onnx` files under `directory`.""" + found: Dict[str, Path] = {} + for m in modules: + hits = sorted(directory.glob(f"{m}*.onnx")) + if len(hits) > 1: + raise SystemExit(f"{directory}: several graphs for module {m!r}: {[h.name for h in hits]}") + if hits: + found[m] = hits[0] + if not found: + raise SystemExit(f"{directory}: no llm*.onnx / dit*.onnx / expert*.onnx plugin graphs") + return found + + def nodes(self, module: str) -> Dict[str, Dict[str, Any]]: + if module not in self._nodes: + import onnx + + if module not in self.files: + raise SystemExit(f"no plugin graph for module {module!r} (have {sorted(self.files)})") + model = onnx.load(str(self.files[module]), load_external_data=True) + plug: Dict[str, Dict[str, Any]] = {} + for node in model.graph.node: + if node.domain: + plug[node.name] = {a.name: self._attr(a) for a in node.attribute} + plug[node.name]["__op__"] = node.op_type + self._nodes[module] = plug + return self._nodes[module] + + @staticmethod + def _attr(a: Any) -> Any: + import onnx + + if a.type == onnx.AttributeProto.STRING: + return a.s + if a.type == onnx.AttributeProto.INT: + return int(a.i) + if a.type == onnx.AttributeProto.FLOAT: + return float(a.f) + if a.type == onnx.AttributeProto.INTS: + return list(a.ints) + return None + + def node(self, module: str, name: str) -> Dict[str, Any]: + n = self.nodes(module).get(name) + if n is None: + raise SystemExit(f"{self.files[module]}: no plugin node {name!r}") + return n + + @staticmethod + def _weight(a: Dict[str, Any], stem: str, n: int, k: int) -> Tuple[np.ndarray, np.ndarray]: + i4, i8 = f"{stem}_i4", f"{stem}_i8" + if i4 in a: + codes = _as_codes(a[i4], n, k, 4) + elif i8 in a: + codes = _as_codes(a[i8], n, k, 8) + else: + raise SystemExit(f"plugin node has neither {i4} nor {i8}") + scale = np.frombuffer(a[f"{stem}_scale"], np.float32) + if scale.shape[0] != n: + raise SystemExit(f"{stem}_scale has {scale.shape[0]} entries, N={n}") + return codes, scale.copy() + + @staticmethod + def _vec(a: Dict[str, Any], key: str, k: int) -> np.ndarray: + v = np.frombuffer(a.get(key, b""), np.float32) + if v.shape[0] != k: + raise SystemExit(f"{key} has {v.shape[0]} entries, K={k} (vla.cpp runs the fold-before SmoothQuant vector only)") + return v.copy() + + def scheme(self, module: str) -> Tuple[str, Dict[str, Any]]: + if module not in self._schemes: + raise SystemExit(f"no scheme recorded for module {module!r}") + return self._schemes[module] + + # LLM + def llm_layers(self) -> int: + return sum(1 for n in self.nodes("llm") if n.startswith("L") and n.endswith("_rmsnorm_qkv")) + + def llm_site(self, i: int, key: str) -> Rec: + a = self.node("llm", LLM_NODES[key].format(i=i)) + codes, scale = self._weight(a, "weight", int(a["N"]), int(a["K"])) + return Rec(codes, scale) + + def llm_gamma(self, i: int, key: str, ckpt_gain: torch.Tensor, *, gemma: bool) -> torch.Tensor: # noqa: ARG002 + a = self.node("llm", LLM_NODES[key].format(i=i)) + return _gamma_from_bytes(a["gamma"], int(a["K"])) + + def llm_rot_block(self) -> int: + return int(self.node("llm", LLM_NODES["qkv"].format(i=0)).get("rot_block_size", 64)) + + # DiT + def _dit_node(self, i: int, kind: str) -> Optional[Dict[str, Any]]: + nodes = self.nodes("dit") + for suffix in ("_int4", "_full", ""): + n = nodes.get(f"block{i}_{kind}{suffix}") + if n is not None: + return n + return None + + def dit_blocks(self) -> int: + n = 0 + while self._dit_node(n, "ffn") is not None: + n += 1 + return n + + def dit_block(self, i: int) -> Dict[str, Rec]: + out: Dict[str, Rec] = {} + ffn = self._dit_node(i, "ffn") + if ffn is None: + raise SystemExit(f"{self.files['dit']}: no ffn node for block {i}") + k, inner = int(ffn["K"]), int(ffn["inner_dim"]) + c, s = self._weight(ffn, "weight_proj0", inner, k) + out["ffn0"] = Rec(c, s, self._vec(ffn, "act_scale_pre0", k)) + c, s = self._weight(ffn, "weight_proj2", k, inner) + out["ffn2"] = Rec(c, s, self._vec(ffn, "act_scale_pre2", inner)) + sa, ca = self._dit_node(i, "selfattn"), self._dit_node(i, "crossattn") + if sa is not None: + d = int(sa["inner_dim"]) + c, s = self._weight(sa, "weight_qkv", 3 * d, k) + out["qkv"] = Rec(c, s, self._vec(sa, "act_scale_pre_in", k)) + c, s = self._weight(sa, "weight_o", k, d) + out["o"] = Rec(c, s, self._vec(sa, "act_scale_pre_o", d)) + elif ca is not None: + d, k_enc = int(ca["inner_dim"]), int(ca["K_enc"]) + c, s = self._weight(ca, "weight_q", d, k) + out["q"] = Rec(c, s, self._vec(ca, "act_scale_pre_in", k)) + nodes = self.nodes("dit") + enc = nodes.get("encoder_prequant_int4") or nodes.get("encoder_prequant") + if enc is None: + raise SystemExit(f"{self.files['dit']}: no encoder_prequant node") + c, s = self._weight(ca, "weight_kv", 2 * d, k_enc) + out["kv"] = Rec(c, s, self._vec(enc, "act_scale_pre_enc", k_enc)) + c, s = self._weight(ca, "weight_o", k, d) + out["o"] = Rec(c, s, self._vec(ca, "act_scale_pre_o", d)) + else: + raise SystemExit(f"{self.files['dit']}: block {i} has neither a selfattn nor a crossattn node") + return out + + def dit_rot_block(self) -> int: + a = self._dit_node(0, "ffn") or {} + return int(a.get("rot_block_size", a.get("block_size", 64))) + + # Gemma expert + def expert_layers(self) -> int: + return sum(1 for n in self.nodes("expert") if n.startswith("G") and "_qkv_plr" in n) + + def expert_site(self, i: int, key: str) -> Rec: + nodes = self.nodes("expert") + a = nodes.get(f"G{i}_{key}_plr4") or nodes.get(f"G{i}_{key}_plr") + if a is None: + raise SystemExit(f"{self.files['expert']}: no plugin node G{i}_{key}_plr4 / _plr") + n, k = int(a["N"]), int(a["K"]) + codes, scale = self._weight(a, "weight", n, k) + return Rec(codes, scale, self._vec(a, "act_scale_pre", k)) + + def provenance(self) -> str: + return f"plugin graphs {self.root}" + + +def open_source(root: Path) -> Source: + if root.is_file() and root.suffix == ".safetensors": + root = root.parent + if (root / FQ_MANIFEST).is_file(): + # Both FoldQuantVLA formats keep their manifest under the same name; `format` tells them apart. + fmt = json.loads((root / FQ_MANIFEST).read_text()).get("format") + if fmt == FQC_FORMAT: + return FoldQuantCheckpoint(root) + if fmt == FQ_FORMAT: + return FoldQuantState(root) + raise SystemExit(f"{root / FQ_MANIFEST}: FoldQuant format {fmt!r} is not one this converter reads " + f"({FQC_FORMAT!r}, {FQ_FORMAT!r})") + raise SystemExit(f"{root}: not a FoldQuantVLA quantized model (no {FQ_MANIFEST})") + + +# --------------------------------------------------------------------------- +# families +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Family: + arch: str + converter: str # vla.cpp converter module with convert(ckpt, out, *, writer_factory, **kwargs) + llm_prefix: str # GGUF block prefix of the LLM + action_module: str # "dit" (GR00T) or "expert" (pi0.5) + action_prefix: str # GGUF block prefix of the action module + gemma: bool # norm gains are applied as 1 + w (pi0.5); the GGUF stores the folded gain - 1 + f32_overrides: bool # folded gains written as F32 (else in the checkpoint's dtype) + converter_kwargs: Dict[str, Any] + + +FAMILIES: Dict[str, Family] = { + "groot_n1_5": Family("gr00t_n1_5", "convert_gr00t_n1_5_to_gguf", "vlm.blk", "dit", "aex.dit", False, False, {}), + # N1.6's tower sees 252 px frames (see convert_gr00t_n1_6_to_gguf.py --vision-size). + "groot_n1_6": Family("gr00t_n1_6", "convert_gr00t_n1_6_to_gguf", "vlm.blk", "dit", "aex.dit", False, False, {"vision_size": 252}), + "groot_n1_7": Family("gr00t_n1_7", "convert_gr00t_n1_7_to_gguf", "vlm.blk", "dit", "aex.dit", False, False, {}), + "pi05": Family("pi05", "convert_pi05_to_gguf", "vlm.blk", "expert", "aex.blk", True, True, {}), +} + +LLM_PROJ = {"qkv": ("attn_q", "attn_k", "attn_v"), "o": ("attn_o",), "gateup": ("ffn_gate", "ffn_up"), "down": ("ffn_down",)} +EXPERT_PROJ = {"qkv": ("attn_q", "attn_k", "attn_v"), "o": ("attn_o",), "gu": ("ffn_gate", "ffn_up"), "dn": ("ffn_down",)} + + +def _check_schemes(src: Source, action_module: str) -> Tuple[str, str]: + ls, _ = src.scheme("llm") + as_, _ = src.scheme(action_module) + if ls not in LLM_SCHEMES: + raise SystemExit(f"LLM scheme {ls!r} has no vla.cpp kernel (supported: {sorted(LLM_SCHEMES)})") + if as_ not in ACTION_SCHEMES: + raise SystemExit(f"action scheme {as_!r} has no vla.cpp kernel (supported: {sorted(ACTION_SCHEMES)}; " + "`*_sr` action schemes use a dense learned rotation, vla.cpp runs the fixed butterfly)") + return ls, as_ + + +def _rows(n: int, kv_dim: int, key: str) -> List[int]: + if key == "qkv": + return [n - 2 * kv_dim, kv_dim, kv_dim] + if key in ("gateup", "gu", "kv"): + return [n // 2, n // 2] + return [n] + + +def build_sites(src: Source, fam: Family, ckpt: Checkpoint, conv: Any) -> Tuple[Dict[str, Site], Dict[str, np.ndarray], QuantKV, Dict[str, Any]]: + """GGUF sites, folded-gain overrides and `.quant.*` metadata for one quantized model.""" + llm_scheme, act_scheme = _check_schemes(src, fam.action_module) + _, llm_cfg = src.scheme("llm") + _, act_cfg = src.scheme(fam.action_module) + sites: Dict[str, Site] = {} + overrides: Dict[str, np.ndarray] = {} + gains: Dict[str, torch.Tensor] = {} + + # -- LLM -------------------------------------------------------------------------- + lm_root = conv.PFX_VLM if fam.gemma else conv.LM_ROOT + n_layers = src.llm_layers() + if n_layers == 0: + raise SystemExit("the quantized model records no LLM layers") + kv_dim = ckpt.shape(f"{lm_root}.layers.0.self_attn.k_proj.weight")[0] + widths: Dict[str, set] = {k: set() for k in LLM_SITES} + for i in range(n_layers): + for key in LLM_SITES: + rec = src.llm_site(i, key) + widths[key].add(rec.bits) + for name, (c, s) in zip(LLM_PROJ[key], rec.split(_rows(rec.n_out, kv_dim, key))): + sites[f"{fam.llm_prefix}.{i}.{name}.weight"] = Site(codes=c, wscale=s, ascale=None, bits=rec.bits) + for key, norm, gguf_norm in (("qkv", "input_layernorm", "attn_norm"), ("gateup", "post_attention_layernorm", "ffn_norm")): + gain = ckpt.tensor(f"{lm_root}.layers.{i}.{norm}.weight") + g = src.llm_gamma(i, key, gain, gemma=fam.gemma) + if g.numel() != gain.numel(): + raise SystemExit(f"L{i}_{key}: folded gain has {g.numel()} entries, the checkpoint's {norm} {gain.numel()}") + gains[f"L{i}_{key}"] = g + # vla.cpp loads Gemma norms as 1 + w, so a Gemma file gets gain - 1 (exact in fp32 from bf16 values). + overrides[f"{fam.llm_prefix}.{i}.{gguf_norm}.weight"] = ((g.double() - 1.0).float().numpy() if fam.gemma + else g.float().numpy()) + for key, ws in widths.items(): + if len(ws) != 1: + raise SystemExit(f"LLM site {key!r} has widths {sorted(ws)} across layers; the GGUF metadata carries one width per site type") + wbits = int(llm_cfg.get("bits", next(iter(widths["qkv"])))) + site_bits = {k: next(iter(ws)) for k, ws in widths.items() if next(iter(ws)) != wbits} + act_bits = int(llm_cfg.get("act_bits") or wbits) + clip = llm_cfg.get("act_clip", llm_cfg.get("act_clip_ratio", 1.0)) + if isinstance(clip, dict): + raise SystemExit("per-site learned activation clips are not representable in the GGUF metadata yet") + rot_bs = int(llm_cfg.get("rot_bs", llm_cfg.get("rot_block_size", src.llm_rot_block()))) + + # -- action module ------------------------------------------------------------------- + act_widths: set = set() + if fam.action_module == "dit": + n_blocks = src.dit_blocks() + if n_blocks == 0: + raise SystemExit("the quantized model records no DiT blocks") + for i in range(n_blocks): + base = f"{fam.action_prefix}.{i}" + d = ckpt.shape(f"{conv.AHK}.model.transformer_blocks.{i}.attn1.to_k.weight")[0] + recs = src.dit_block(i) + if "qkv" in recs: + r = recs["qkv"] + for name, (c, s) in zip(("attn_q", "attn_k", "attn_v"), r.split([r.n_out - 2 * d, d, d])): + sites[f"{base}.{name}.weight"] = Site(codes=c, wscale=s, ascale=r.ascale, bits=r.bits) + act_widths.add(r.bits) + else: + r = recs["q"] + sites[f"{base}.attn_q.weight"] = Site(codes=r.codes, wscale=r.scale, ascale=r.ascale, bits=r.bits) + act_widths.add(r.bits) + r = recs["kv"] + for name, (c, s) in zip(("attn_k", "attn_v"), r.split([d, d])): + sites[f"{base}.{name}.weight"] = Site(codes=c, wscale=s, ascale=r.ascale, bits=r.bits) + act_widths.add(r.bits) + for key, name in (("o", "attn_o"), ("ffn0", "ff0"), ("ffn2", "ff2")): + r = recs[key] + sites[f"{base}.{name}.weight"] = Site(codes=r.codes, wscale=r.scale, ascale=r.ascale, bits=r.bits) + act_widths.add(r.bits) + act_rot = int(act_cfg.get("rot_bs", act_cfg.get("rot_block_size", src.dit_rot_block()))) + else: + n_exp = src.expert_layers() + if n_exp == 0: + raise SystemExit("the quantized model records no expert layers") + exp_kv = ckpt.shape(f"{conv.PFX_AEX}.layers.0.self_attn.k_proj.weight")[0] + for i in range(n_exp): + for key in EXPERT_SITES: + r = src.expert_site(i, key) + for name, (c, s) in zip(EXPERT_PROJ[key], r.split(_rows(r.n_out, exp_kv, key))): + sites[f"{fam.action_prefix}.{i}.{name}.weight"] = Site(codes=c, wscale=s, ascale=r.ascale, bits=r.bits) + act_widths.add(r.bits) + act_rot = int(act_cfg.get("rot_bs", act_cfg.get("rot_block_size", 64))) + if len(act_widths) != 1: + raise SystemExit(f"action sites have mixed widths {sorted(act_widths)}; vla.cpp runs one width per action module") + abits = act_widths.pop() + fold_order = str(act_cfg.get("sq_fold_order", (act_cfg.get("params") or {}).get("sq_fold_order", "before"))) + if fold_order != "before": + raise SystemExit(f"action sq_fold_order {fold_order!r}: vla.cpp divides the SmoothQuant vector in before the rotation only") + + kv: QuantKV = { + "method": ("str", "foldquant"), + "applied_at": ("str", src.producer), + "provenance": ("str", src.provenance()), + "scheme_llm": ("str", llm_scheme), + "llm_weight_bits": ("u32", wbits), + "llm_act_bits": ("u32", act_bits), + "llm_rot_block_size": ("u32", rot_bs), + "act_clip_ratio": ("f32", float(clip)), + "site_bits": ("str", ",".join(f"{k}:{v}" for k, v in site_bits.items())), + "scheme_action": ("str", act_scheme), + "action_weight_bits": ("u32", abits), + "action_act_bits": ("u32", abits), + "action_rot_block_size": ("u32", act_rot), + "action_fold_order": ("str", fold_order), + } + return sites, overrides, kv, {"gains": gains} + + +def check_against_graphs(src: Source, fam: Family, ckpt: Checkpoint, conv: Any, sites: Dict[str, Site], + gains: Dict[str, torch.Tensor], onnx_dir: Path) -> int: + """Byte-compare the sites and folded gains against the plugin graphs in `onnx_dir`. Returns the site count.""" + graphs = PluginGraphs(PluginGraphs.find(onnx_dir), family=src.family, + schemes={m: src.scheme(m) for m in ("llm", fam.action_module)}) + ref_sites, _, _, ref_extra = build_sites(graphs, fam, ckpt, conv) + bad: List[str] = [] + if set(ref_sites) != set(sites): + bad.append(f"site sets differ ({len(sites)} vs {len(ref_sites)} in the graphs)") + for name in sorted(set(ref_sites) & set(sites)): + a, b = sites[name], ref_sites[name] + if a.bits != b.bits or a.codes.tobytes() != b.codes.tobytes(): + bad.append(f"{name} codes") + if a.wscale.astype(np.float32).tobytes() != b.wscale.astype(np.float32).tobytes(): + bad.append(f"{name} wscale") + if (a.ascale is None) != (b.ascale is None) or (a.ascale is not None and a.ascale.tobytes() != b.ascale.tobytes()): + bad.append(f"{name} ascale") + for key, g in gains.items(): + want = ref_extra["gains"].get(key) + if want is None or g.dtype != want.dtype or not torch.equal(g, want): + bad.append(f"{key} gamma") + if bad: + raise SystemExit(f"--check-onnx: {len(bad)} mismatch(es) against {onnx_dir}: {bad[:8]}") + return len(ref_sites) + + +# --------------------------------------------------------------------------- +# converter arguments per family +# --------------------------------------------------------------------------- + + +def converter_kwargs(src: Source, fam: Family, ckpt: Checkpoint, args: argparse.Namespace) -> Dict[str, Any]: + kwargs: Dict[str, Any] = dict(fam.converter_kwargs) + if fam.arch != "pi05": + return kwargs + stats = args.dataset_stats + if stats is None: + hits = sorted(ckpt.root.glob("assets/*/*/norm_stats.json")) + if len(hits) != 1: + raise SystemExit(f"{ckpt.root}: expected one assets/*/*/norm_stats.json, found {len(hits)}; pass --dataset-stats") + stats = hits[0] + kwargs["dataset_stats"] = Path(stats) + cfg_path = ckpt.root / "config.json" + cfg_json = json.loads(cfg_path.read_text()) if cfg_path.is_file() else {} + if cfg_json.get("type") != "pi05": + if args.config_json: + kwargs["config"] = json.loads(Path(args.config_json).read_text()) + else: + meta = src.export_metadata() if isinstance(src, FoldQuantState) else {} # also FoldQuantCheckpoint + norm = json.loads(Path(stats).read_text()).get("norm_stats", {}) + horizon = int(cfg_json.get("action_horizon", meta.get("action_horizon", 50))) + dim = int(cfg_json.get("action_dim", meta.get("action_dim", 32))) + kwargs["config"] = { + "chunk_size": horizon, "n_action_steps": horizon, + "num_inference_steps": int(meta.get("num_steps", 10)), + "max_state_dim": dim, "max_action_dim": dim, + "min_period": 4e-3, "max_period": 4.0, + "tokenizer_max_length": int(cfg_json.get("max_token_len", 200)), + "input_features": {"observation.state": {"shape": [len(norm["state"]["mean"])]}}, + "output_features": {"action": {"shape": [len(norm["actions"]["mean"])]}}, + } + return kwargs + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--quantized-model", type=Path, required=True, + help="FoldQuantVLA quantized model dir (foldquant_quant.json): a quantized checkpoint or an earlier fake-quant state") + ap.add_argument("--out", type=Path, required=True, help="output GGUF path") + ap.add_argument("--ckpt", type=Path, default=None, + help="[fake-quant state] base checkpoint dir (default: the state dir itself); a quantized checkpoint is its own base") + ap.add_argument("--check-onnx", type=Path, default=None, + help="FoldQuantVLA plugin ONNX dir (llm_bf16.onnx, dit_bf16.onnx / expert_bf16.onnx) to byte-compare every site against") + ap.add_argument("--dataset-stats", type=Path, default=None, help="[pi05] normalizer stats (default: the checkpoint's assets/*/*/norm_stats.json)") + ap.add_argument("--config-json", type=Path, default=None, help="[pi05] lerobot policy fields for an OpenPI checkpoint (default: from the export metadata)") + args = ap.parse_args() + + src = open_source(args.quantized_model) + fam = FAMILIES.get(src.family) + if fam is None: + raise SystemExit(f"family {src.family!r} has no vla.cpp FoldQuant sites (supported: {sorted(FAMILIES)})") + if isinstance(src, FoldQuantCheckpoint) and args.ckpt and Path(args.ckpt).resolve() != src.root: + raise SystemExit("a FoldQuantVLA quantized checkpoint is its own base checkpoint (configs, assets and every " + "unquantized weight are in it); drop --ckpt") + ckpt_root = args.ckpt or src.base_checkpoint + if ckpt_root is None: + raise SystemExit("the quantized model records no checkpoint path; pass --ckpt ") + ckpt = Checkpoint(Path(ckpt_root)) + conv = importlib.import_module(fam.converter) + + sites, overrides, kv, extra = build_sites(src, fam, ckpt, conv) + if isinstance(src, FoldQuantCheckpoint): + src.check_accounting() + if src.float_projections(): + print(f"{len(src.float_projections())} DiT adaLN projection(s) are INT4 in the checkpoint; vla.cpp runs them in " + "bf16 from the dequantized weight (codes x bf16 scale, as the TensorRT AdaLNModInt4 plugin uses it)") + print(f"{src.producer} {src.family}: {len(sites)} sites, {len(overrides)} folded norm gains, " + f"llm={kv['scheme_llm'][1]} W{kv['llm_weight_bits'][1]}A{kv['llm_act_bits'][1]}" + f"{' site_bits=' + kv['site_bits'][1] if kv['site_bits'][1] else ''} " + f"action={kv['scheme_action'][1]} W{kv['action_weight_bits'][1]}A{kv['action_act_bits'][1]}") + if args.check_onnx: + n = check_against_graphs(src, fam, ckpt, conv, sites, extra["gains"], args.check_onnx) + print(f"--check-onnx: {n} sites and every folded gain byte-identical to {args.check_onnx}") + + factory, holder = quantizing_writer_factory(sites, overrides, kv, f32_overrides=fam.f32_overrides) + with base_view(conv, ckpt): + out = conv.convert(ckpt.root, args.out, writer_factory=factory, **converter_kwargs(src, fam, ckpt, args)) + missing = holder[0].unconsumed() if holder else [""] + if missing: + args.out.unlink(missing_ok=True) + raise SystemExit(f"{len(missing)} site(s)/override(s) were never written (name mismatch): {missing[:8]}") + print(f"done: {out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/foldquant_dequant.py b/scripts/foldquant_dequant.py new file mode 100644 index 0000000..b343ecc --- /dev/null +++ b/scripts/foldquant_dequant.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Turn FoldQuant sites of a GGUF back into BF16 float weights (a debugging aid). + +A FoldQuant site computes ``y = Wq . q(H (x / s))`` with ``Wq`` the INT8 codes +times the per-row scale, ``H`` the block Hadamard (symmetric, orthonormal) and +``s`` the optional SmoothQuant vector. Dropping the activation quantizer gives +the float weight ``W_eff = Wq . H . diag(1/s)`` that reproduces the site up to +quantization noise, with every fold (SQ in the norm gammas, rotation) kept. +Running the result on a float path isolates the FoldQuant kernels from the +rest of the graph: a bisection over ``--modules`` says which module's +integer path deviates from the reference. + + python scripts/foldquant_dequant.py --in int8.gguf --out deq.gguf [--modules llm,action] +""" +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import gguf # noqa: E402 + +import foldquant_ref as fq # noqa: E402 +from quantize_gguf import to_f32 # noqa: E402 + +LLM_SITE = re.compile(r"^vlm\.blk\.\d+\.(attn_q|attn_k|attn_v|attn_o|ffn_gate|ffn_up|ffn_down)\.weight$") +DIT_SITE = re.compile(r"^aex\.dit\.\d+\.(attn_q|attn_k|attn_v|attn_o|ff0|ff2)\.weight$") + + +def _bf16(a: np.ndarray) -> np.ndarray: + """Round-to-nearest-even float32 -> bf16 bit pattern (uint16).""" + u = np.ascontiguousarray(a, dtype=np.float32).view(np.uint32) + return ((u + 0x7FFF + ((u >> 16) & 1)) >> 16).astype(np.uint16) + + +def _selfcheck() -> None: + rng = np.random.default_rng(0) + w = rng.standard_normal((128, 256)).astype(np.float32) + s = rng.uniform(0.5, 2.0, 256).astype(np.float32) + folded = fq.fold_weight(w * s[None, :], 64) # what the exporter stores (SQ then rotate) + back = fq.fwht_rows(folded, 64) / s[None, :] + assert np.allclose(back, w, atol=1e-4), "fold/dequant are not inverses; check the Hadamard convention" + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--in", dest="src", required=True, type=Path) + ap.add_argument("--out", dest="dst", required=True, type=Path) + ap.add_argument("--modules", default="llm,action", help="comma list of llm, action: which sites to dequantize") + args = ap.parse_args() + modules = {m.strip() for m in args.modules.split(",") if m.strip()} + _selfcheck() + + r = gguf.GGUFReader(str(args.src)) + arch = r.fields["general.architecture"].contents() + qp = f"{arch}.quant." + if not any(k.startswith(qp) for k in r.fields): + raise SystemExit(f"{args.src} is not a FoldQuant GGUF") + tensors = {t.name: t for t in r.tensors} + rot = { + "llm": int(r.fields[qp + "llm_rot_block_size"].contents()), + "action": int(r.fields[qp + "action_rot_block_size"].contents()), + } + fold_before = str(r.fields[qp + "action_fold_order"].contents()) == "before" + + w = gguf.GGUFWriter(str(args.dst), arch) + meta = {"GGUF.version", "GGUF.tensor_count", "GGUF.kv_count", "general.architecture"} + for name, f in r.fields.items(): + if name in meta: + continue + val = f.contents() + if name == qp + "provenance": + val = f"{val}; dequantized {sorted(modules)} by scripts/foldquant_dequant.py" + if f.types and f.types[0] == gguf.GGUFValueType.ARRAY: + w.add_array(name, val) + else: + w.add_key_value(name, val, f.types[0]) + + n_deq = 0 + skip: set[str] = set() + for t in r.tensors: + if t.name in skip: + continue + mod = "llm" if LLM_SITE.match(t.name) else ("action" if DIT_SITE.match(t.name) else None) + if mod in modules and t.tensor_type == gguf.GGMLQuantizationType.I8: + base = t.name[: -len(".weight")] + codes = np.ascontiguousarray(t.data).view(np.int8) + n_out, kpack = int(t.shape[1]), int(t.shape[0]) # reader shape is ggml ne order + if int(r.fields[qp + f"{mod}_weight_bits"].contents()) == 4: + codes = fq.unpack_nibbles(codes.reshape(n_out, kpack)) + k_in = codes.shape[1] if codes.ndim == 2 else kpack + codes = codes.reshape(n_out, k_in).astype(np.float32) + wscale = np.asarray(tensors[base + ".wscale"].data, dtype=np.float32).reshape(n_out) + wq = codes * wscale[:, None] + bs = fq.rotation_block_for(k_in, rot[mod]) + w_eff = fq.fwht_rows(wq, bs) if bs > 1 else wq + asc_t = tensors.get(base + ".ascale") + if asc_t is not None: + s = np.asarray(asc_t.data, dtype=np.float32).reshape(k_in) + # fold-before: y = Wq H (x/s) -> W_eff = (Wq H) / s + # fold-after : y = Wq ((H x)/s) -> W_eff = (Wq diag(1/s)) H + w_eff = w_eff / s[None, :] if fold_before else (fq.fwht_rows(wq / s[None, :], bs) if bs > 1 else wq / s[None, :]) + skip.add(base + ".ascale") + skip.add(base + ".wscale") + w.add_tensor(t.name, _bf16(w_eff), raw_shape=[n_out, k_in], raw_dtype=gguf.GGMLQuantizationType.BF16) + n_deq += 1 + continue + data = np.ascontiguousarray(t.data) + if t.tensor_type == gguf.GGMLQuantizationType.BF16: + data = data.view(np.uint16) + elif t.tensor_type == gguf.GGMLQuantizationType.F32: + data = data.astype(np.float32, copy=False) + w.add_tensor(t.name, data, raw_dtype=t.tensor_type) + + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + print(f"dequantized {n_deq} sites ({sorted(modules)}) -> {args.dst}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/foldquant_fake_export.py b/scripts/foldquant_fake_export.py new file mode 100644 index 0000000..d721cff --- /dev/null +++ b/scripts/foldquant_fake_export.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Turn a float GR00T GGUF into a contract-conformant FoldQuant GGUF without +calibration: block-Hadamard rotation of every site's weight plus per-row +symmetric INT8 (or nibble-packed INT4) rounding, no SmoothQuant (so no +.ascale and unchanged norms). It is the in-repo producer of the format in +docs/QUANTIZATION.md, for kernel bring-up, benchmarks and CI without a +calibrated model; a calibrated file, with the SmoothQuant folds, comes from a +FoldQuantVLA quantized model through scripts/convert_quantized_model_to_gguf.py. + + python scripts/foldquant_fake_export.py --in n17-bf16.gguf --out n17-fq.gguf + [--rot-block 64] [--wbits 8|4] [--modules llm,action] +""" + +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path + +import numpy as np +import gguf + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import foldquant_ref as fq # noqa: E402 +from quantize_gguf import to_f32 # noqa: E402 + +LLM_SITE = re.compile(r"^vlm\.blk\.\d+\.(attn_q|attn_k|attn_v|attn_o|ffn_gate|ffn_up|ffn_down)\.weight$") +DIT_SITE = re.compile(r"^aex\.dit\.\d+\.(attn_q|attn_k|attn_v|attn_o|ff0|ff2)\.weight$") +# The scheme label this file records (one per weight width): rotation plus +# round-to-nearest, no SmoothQuant, no calibration. A label only; the runtime +# reads the widths, not the name. +SCHEME = {8: "foldquant_fake_w8a8", 4: "foldquant_fake_w4a4"} + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--in", dest="src", required=True, type=Path) + ap.add_argument("--out", dest="dst", required=True, type=Path) + ap.add_argument("--rot-block", type=int, default=64, help="nominal Hadamard block (narrowed per site)") + ap.add_argument("--wbits", type=int, default=8, choices=(8, 4)) + ap.add_argument("--abits", type=int, default=8, choices=(8, 4), help="activation width recorded in the metadata") + ap.add_argument("--modules", default="llm,action", help="comma list of llm, action") + ap.add_argument("--sites", default=None, help="regex; only matching site names are quantized (bisection aid)") + args = ap.parse_args() + modules = {m.strip() for m in args.modules.split(",") if m.strip()} + site_filter = re.compile(args.sites) if args.sites else None + + r = gguf.GGUFReader(str(args.src)) + arch = r.fields["general.architecture"].contents() + quant_prefix = f"{arch}.quant." + if any(k.startswith(quant_prefix) for k in r.fields): + raise SystemExit(f"{args.src} is already a FoldQuant GGUF") + + w = gguf.GGUFWriter(str(args.dst), arch) + meta = {"GGUF.version", "GGUF.tensor_count", "GGUF.kv_count", "general.architecture"} + for name, f in r.fields.items(): + if name in meta: + continue + if f.types and f.types[0] == gguf.GGUFValueType.ARRAY: + w.add_array(name, f.contents()) + else: + w.add_key_value(name, f.contents(), f.types[0]) + + kv = lambda k: quant_prefix + k # noqa: E731 + w.add_string(kv("method"), "foldquant") + w.add_string(kv("applied_at"), "foldquant_fake_export") + w.add_string(kv("provenance"), f"scripts/foldquant_fake_export.py rot_block={args.rot_block} wbits={args.wbits} (no calibration, no SmoothQuant)") + for mod in ("llm", "action"): + on = mod in modules + w.add_string(kv(f"scheme_{mod}"), SCHEME[args.wbits] if on else "float") + w.add_uint32(kv(f"{mod}_weight_bits"), args.wbits if on else 16) + w.add_uint32(kv(f"{mod}_act_bits"), args.abits if on else 32) + w.add_uint32(kv(f"{mod}_rot_block_size"), args.rot_block if on else 0) + w.add_string(kv("action_fold_order"), "before") + w.add_float32(kv("act_clip_ratio"), 1.0) + w.add_string(kv("site_bits"), "") + + n_sites = skipped = 0 + bytes_in = bytes_out = 0 + for t in r.tensors: + src_bytes = int(t.data.nbytes) + bytes_in += src_bytes + is_llm = bool(LLM_SITE.match(t.name)) and "llm" in modules + is_dit = bool(DIT_SITE.match(t.name)) and "action" in modules + if site_filter is not None and not site_filter.search(t.name): + is_llm = is_dit = False + f32 = to_f32(t) if (is_llm or is_dit) else None + if f32 is not None and f32.ndim == 2 and f32.shape[0] % 64 == 0 and f32.shape[1] % 64 == 0: + n_out, k_in = f32.shape + bs = fq.rotation_block_for(k_in, args.rot_block) + folded = fq.fold_weight(f32, bs) + codes, wscale = fq.weight_quant_per_row(folded, args.wbits) + body = fq.pack_nibbles(codes) if args.wbits == 4 else codes.view(np.uint8) + w.add_tensor(t.name, np.ascontiguousarray(body), raw_shape=[n_out, body.shape[1]], + raw_dtype=gguf.GGMLQuantizationType.I8) + base = t.name[: -len(".weight")] + w.add_tensor(base + ".wscale", np.ascontiguousarray(wscale, dtype=np.float32), + raw_shape=[n_out], raw_dtype=gguf.GGMLQuantizationType.F32) + bytes_out += int(body.nbytes) + int(wscale.nbytes) + n_sites += 1 + continue + if is_llm or is_dit: + skipped += 1 + print(f" note: {t.name} shape {tuple(int(s) for s in t.shape)} not a multiple of 64, kept float") + data = np.ascontiguousarray(t.data) + if t.tensor_type == gguf.GGMLQuantizationType.BF16: + data = data.view(np.uint16) + elif t.tensor_type == gguf.GGMLQuantizationType.F32: + data = data.astype(np.float32, copy=False) + w.add_tensor(t.name, data, raw_dtype=t.tensor_type) + bytes_out += src_bytes + + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + print(f"foldquant W{args.wbits}: {n_sites} sites quantized ({skipped} kept float), " + f"weights {bytes_in/1e9:.2f} GB -> {bytes_out/1e9:.2f} GB ({100*bytes_out/max(bytes_in,1):.0f}%)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/foldquant_ref.py b/scripts/foldquant_ref.py new file mode 100644 index 0000000..9ec1f8e --- /dev/null +++ b/scripts/foldquant_ref.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""numpy reference for the FoldQuant contract (docs/QUANTIZATION.md). + +Mirrors src/foldquant_ref.h operation for operation in float32 so that, for the +sites without a fused RMSNorm, the INT8 codes it produces are bit-identical to +the runtime's (the norm's sum of squares is the one reduction whose order numpy +cannot reproduce; there the runtime may differ by one code at a rounding +boundary). tests/py/test_foldquant_ref.py pins this against the C++ test's +golden checksum. numpy only, so it also serves scripts/foldquant_fake_export.py. +""" + +from __future__ import annotations + +import numpy as np + +QMAX = {8: 127.0, 4: 7.0} +ACT_TAIL = 16 # bytes after the codes in an activation-blob row +SCALE_FLOOR = 1e-12 + + +def lcg(seed: int, n: int) -> np.ndarray: + """Same generator as tests/test_foldquant_cpu_op.cpp: n floats in [-1, 1].""" + # Every operation in float32, like the C++: a float64 division rounded to + # float32 afterwards differs by an ulp often enough to flip a code. + out = np.empty(n, dtype=np.float32) + s = seed & 0xFFFFFFFF + two, one, denom = np.float32(2.0), np.float32(1.0), np.float32(65535.0) + for i in range(n): + s = (s * 1664525 + 1013904223) & 0xFFFFFFFF + out[i] = (np.float32((s >> 8) & 0xFFFF) / denom) * two - one + return out + + +def rotation_block_for(k_in: int, nominal: int) -> int: + """Largest power of two <= nominal dividing k_in; 1 means no rotation.""" + bs = nominal + while bs > 1 and k_in % bs: + bs //= 2 + return bs if bs >= 2 else 1 + + +def inv_sqrt_block(bs: int) -> np.float32: + return np.float32(1.0 / np.sqrt(float(bs))) + + +def hadamard(bs: int) -> np.ndarray: + """Normalised natural-order Sylvester Hadamard, float64.""" + h = np.array([[1.0]]) + while h.shape[0] < bs: + h = np.block([[h, h], [h, -h]]) + return h / np.sqrt(bs) + + +def fwht_rows(x: np.ndarray, bs: int) -> np.ndarray: + """In-place-order butterfly on every bs-block of the last axis, float32, + stages h = 1, 2, 4, ... with (a+b, a-b), then * 1/sqrt(bs).""" + y = np.ascontiguousarray(x, dtype=np.float32).copy() + if bs <= 1: + return y + lead = y.shape[:-1] + k = y.shape[-1] + v = y.reshape(*lead, k // bs, bs) + h = 1 + while h < bs: + v = v.reshape(*lead, k // bs, bs // (2 * h), 2, h) + a = v[..., 0, :].copy() + b = v[..., 1, :].copy() + v[..., 0, :] = a + b + v[..., 1, :] = a - b + v = v.reshape(*lead, k // bs, bs) + h *= 2 + v = v * inv_sqrt_block(bs) + return v.reshape(*lead, k).astype(np.float32) + + +def act_quant(x: np.ndarray, *, bits: int = 8, rot_block: int = 64, ascale: np.ndarray | None = None, + fold_before: bool = False, clip: float = 1.0, gamma: np.ndarray | None = None, + eps: float = 1e-6): + """x [T, K] float32 -> (codes int8 [T, K], scales float32 [T]). + + Runtime order: [RMSNorm(gamma)] -> [/ascale if before] -> FWHT -> [/ascale + if after] -> per-token symmetric quant, scale = max(clip*amax/qmax, 1e-12), + q = rint(y * (1/scale)) half to even, clamp +-qmax.""" + y = np.ascontiguousarray(x, dtype=np.float32) + if gamma is not None: + rstd = np.float32(1.0) / np.sqrt((y * y).sum(axis=-1, keepdims=True, dtype=np.float32) / np.float32(y.shape[-1]) + np.float32(eps)) + y = (y * rstd) * gamma.astype(np.float32) + if ascale is not None and fold_before: + y = y / ascale.astype(np.float32) + y = fwht_rows(y, rot_block) + if ascale is not None and not fold_before: + y = y / ascale.astype(np.float32) + amax = np.abs(y).max(axis=-1, keepdims=True).astype(np.float32) + qmax = np.float32(QMAX[bits]) + scale = (np.float32(clip) * amax) / qmax + scale = np.maximum(scale, np.float32(SCALE_FLOOR)).astype(np.float32) + inv = (np.float32(1.0) / scale).astype(np.float32) # reciprocal, like the TensorRT kernels + q = np.clip(np.rint(y * inv), -qmax, qmax).astype(np.int8) + return q, scale.reshape(-1) + + +def weight_quant_per_row(w: np.ndarray, bits: int = 8): + """w [N, K] float32 -> (codes int8 [N, K], wscale float32 [N]); symmetric + per output row, scale = amax/qmax (floored), clamp +-qmax.""" + w = np.ascontiguousarray(w, dtype=np.float32) + qmax = np.float32(QMAX[bits]) + amax = np.abs(w).max(axis=1, keepdims=True).astype(np.float32) + scale = np.maximum(amax / qmax, np.float32(SCALE_FLOOR)).astype(np.float32) + codes = np.clip(np.rint(w / scale), -qmax, qmax).astype(np.int8) + return codes, scale.reshape(-1) + + +def fold_weight(w: np.ndarray, rot_block: int) -> np.ndarray: + """W' = W . H_block^T along K (H is symmetric): the offline rotation whose + inverse the runtime butterfly applies to the activation.""" + if rot_block <= 1: + return np.ascontiguousarray(w, dtype=np.float32) + n, k = w.shape + h = hadamard(rot_block) + v = np.asarray(w, dtype=np.float64).reshape(n, k // rot_block, rot_block) @ h.T + return v.reshape(n, k).astype(np.float32) + + +def pack_nibbles(codes: np.ndarray) -> np.ndarray: + """[.., K] int8 in [-7, 7] -> [.., K/2] uint8, low nibble = even column.""" + c = np.asarray(codes, dtype=np.int8) + lo = (c[..., 0::2].astype(np.uint8)) & 0xF + hi = (c[..., 1::2].astype(np.uint8)) & 0xF + return (lo | (hi << 4)).astype(np.uint8) + + +def unpack_nibbles(packed: np.ndarray) -> np.ndarray: + p = np.asarray(packed, dtype=np.uint8) + lo = ((p << 4).astype(np.int8) >> 4).astype(np.int8) + hi = (p.astype(np.int8) >> 4).astype(np.int8) + out = np.empty(p.shape[:-1] + (p.shape[-1] * 2,), dtype=np.int8) + out[..., 0::2] = lo + out[..., 1::2] = hi + return out + + +def act_blob(codes: np.ndarray, scales: np.ndarray, bits: int = 8) -> np.ndarray: + """The runtime's activation blob: per row, codes (nibble-packed for 4-bit) + then the float32 scale, in a row of K_pack + ACT_TAIL bytes.""" + t, k = codes.shape + body = pack_nibbles(codes) if bits == 4 else codes.astype(np.int8).view(np.uint8) + kp = body.shape[1] + blob = np.zeros((t, kp + ACT_TAIL), dtype=np.uint8) + blob[:, :kp] = body + blob[:, kp:kp + 4] = np.asarray(scales, dtype=np.float32).reshape(t, 1).view(np.uint8) + return blob + + +def gemm_ref(w_codes: np.ndarray, wscale: np.ndarray, x_codes: np.ndarray, xscale: np.ndarray, + bias: np.ndarray | None = None) -> np.ndarray: + """y [T, N] = ((int32 acc) * xscale[t]) * wscale[n] (+ bias[n]) in float32.""" + acc = x_codes.astype(np.int32) @ w_codes.astype(np.int32).T + y = (acc.astype(np.float32) * xscale.astype(np.float32).reshape(-1, 1)) * wscale.astype(np.float32).reshape(1, -1) + if bias is not None: + y = y + bias.astype(np.float32).reshape(1, -1) + return y + + +def fnv1a(data: bytes) -> int: + h = 0xCBF29CE484222325 + for b in data: + h ^= b + h = (h * 0x100000001B3) & 0xFFFFFFFFFFFFFFFF + return h diff --git a/scripts/gguf_quant_writer.py b/scripts/gguf_quant_writer.py new file mode 100644 index 0000000..ed6a9c6 --- /dev/null +++ b/scripts/gguf_quant_writer.py @@ -0,0 +1,139 @@ +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A gguf.GGUFWriter that writes FoldQuant sites (docs/QUANTIZATION.md) as a converter adds them. + +The model converters (`convert__to_gguf.py::convert`) take a +`writer_factory(out, arch)`; hand them `quantizing_writer_factory(...)` and each +tensor named in `sites` is written as integer codes under its own `.weight` +name plus the `.wscale` / `.ascale` sidecars, each name in `overrides` is +replaced by the given array (the SmoothQuant-folded norm gains), and the +`.quant.*` metadata is added. Everything else is the converter's file. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, Optional, Tuple + +import numpy as np + +import gguf + +SITE_ALIGNMENT = 64 + + +@dataclass(frozen=True) +class Site: + """One FoldQuant site as the GGUF carries it. + + codes: int8 (N, K), or uint8 (N, K/2) nibble-packed (low nibble = even column). + wscale: float32 (N,) per-output-row weight scale. + ascale: float32 (K,) SmoothQuant vector divided in before the rotation, or None. + bits: weight width, 4 or 8. + """ + + codes: np.ndarray + wscale: np.ndarray + ascale: Optional[np.ndarray] + bits: int + + @property + def n_out(self) -> int: + return int(self.codes.shape[0]) + + @property + def k_in(self) -> int: + return int(self.codes.shape[1]) * (2 if self.bits == 4 else 1) + + +#: `.quant.` -> ("str" | "u32" | "f32", value) +QuantKV = Dict[str, Tuple[str, Any]] + + +class QuantizingGGUFWriter(gguf.GGUFWriter): + def __init__(self, path: Any, arch: str, *, sites: Dict[str, Site], overrides: Dict[str, np.ndarray], + quant_kv: QuantKV, f32_overrides: bool = False) -> None: + super().__init__(str(path), arch=arch) + self.add_string(f"{arch}.architecture", arch) # the same first key open_writer() writes + self._sites = dict(sites) + self._overrides = dict(overrides) + self._f32_overrides = f32_overrides + self._consumed: set = set() + for key, (kind, value) in quant_kv.items(): + full = f"{arch}.quant.{key}" + if kind == "str": + self.add_string(full, str(value)) + elif kind == "u32": + self.add_uint32(full, int(value)) + elif kind == "f32": + self.add_float32(full, float(value)) + else: + raise ValueError(f"unknown quant KV kind {kind!r} for {full}") + + def add_tensor(self, name: str, tensor: Any, raw_shape: Any = None, raw_dtype: Any = None) -> None: + site = self._sites.get(name) + if site is not None: + shape = list(raw_shape) if raw_shape is not None else list(tensor.shape) + if len(shape) != 2 or int(shape[0]) != site.n_out or int(shape[1]) != site.k_in: + raise ValueError(f"{name}: the converter writes {shape}, the site is (N={site.n_out}, K={site.k_in})") + if site.k_in % SITE_ALIGNMENT or site.n_out % SITE_ALIGNMENT: + raise ValueError(f"{name}: K={site.k_in} N={site.n_out} must be multiples of {SITE_ALIGNMENT}") + base = name[: -len(".weight")] + body = np.ascontiguousarray(site.codes).view(np.uint8) + super().add_tensor(name, body, raw_shape=[site.n_out, int(site.codes.shape[1])], + raw_dtype=gguf.GGMLQuantizationType.I8) + super().add_tensor(base + ".wscale", np.ascontiguousarray(site.wscale, dtype=np.float32), + raw_shape=[site.n_out], raw_dtype=gguf.GGMLQuantizationType.F32) + if site.ascale is not None: + if int(site.ascale.shape[0]) != site.k_in: + raise ValueError(f"{name}: ascale has {site.ascale.shape[0]} entries, K={site.k_in}") + super().add_tensor(base + ".ascale", np.ascontiguousarray(site.ascale, dtype=np.float32), + raw_shape=[site.k_in], raw_dtype=gguf.GGMLQuantizationType.F32) + self._consumed.add(name) + return + override = self._overrides.get(name) + if override is not None: + self._consumed.add(name) + ov = np.ascontiguousarray(override, dtype=np.float32) + if raw_dtype == gguf.GGMLQuantizationType.BF16 and not self._f32_overrides: + import torch + + data = torch.from_numpy(ov).to(torch.bfloat16).view(torch.uint16).numpy() + super().add_tensor(name, data, raw_shape=list(ov.shape), raw_dtype=gguf.GGMLQuantizationType.BF16) + else: + super().add_tensor(name, ov, raw_dtype=gguf.GGMLQuantizationType.F32) + return + super().add_tensor(name, tensor, raw_shape=raw_shape, raw_dtype=raw_dtype) + + def unconsumed(self) -> list: + """Sites / overrides the converter never wrote: a name mismatch, so an error.""" + return sorted((set(self._sites) | set(self._overrides)) - self._consumed) + + +def quantizing_writer_factory(sites: Dict[str, Site], overrides: Dict[str, np.ndarray], quant_kv: QuantKV, *, + f32_overrides: bool = False) -> Tuple[Callable[[Path, str], QuantizingGGUFWriter], list]: + """(factory, holder): pass `factory` as a converter's writer_factory; `holder[0]` is the writer afterwards.""" + holder: list = [] + + def factory(out: Path, arch: str) -> QuantizingGGUFWriter: + Path(out).parent.mkdir(parents=True, exist_ok=True) + print(f"writing {out} ({len(sites)} FoldQuant sites)") + w = QuantizingGGUFWriter(out, arch, sites=sites, overrides=overrides, quant_kv=quant_kv, + f32_overrides=f32_overrides) + holder.append(w) + return w + + return factory, holder diff --git a/scripts/inspect_gguf_quant.py b/scripts/inspect_gguf_quant.py new file mode 100644 index 0000000..db80635 --- /dev/null +++ b/scripts/inspect_gguf_quant.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Print a GGUF's FoldQuant metadata and per-site layout, and check it against +the contract in docs/QUANTIZATION.md: every INT8 weight has an F32 .wscale of +N entries, any .ascale has K entries, K and N are multiples of 64, and fused +groups (q/k/v, k/v) agree on their .ascale. Exit status 1 on a violation, so +an exporter can use it as its conformance gate. + + python scripts/inspect_gguf_quant.py model.gguf [--quiet] +""" + +from __future__ import annotations + +import argparse +import re +import sys +from collections import defaultdict +from pathlib import Path + +import numpy as np +import gguf + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import foldquant_ref as fq # noqa: E402 + +GROUP = re.compile(r"^(.*\.\d+)\.attn_([qkv])$") + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("gguf", type=Path) + ap.add_argument("--quiet", action="store_true", help="only report violations and the summary") + args = ap.parse_args() + + r = gguf.GGUFReader(str(args.gguf)) + arch = r.fields["general.architecture"].contents() + prefix = f"{arch}.quant." + quant = {k[len(prefix):]: f.contents() for k, f in r.fields.items() if k.startswith(prefix)} + tensors = {t.name: t for t in r.tensors} + + if not quant: + print(f"{args.gguf}: no {prefix}* keys (not a FoldQuant GGUF)") + return 0 + print(f"{args.gguf}: arch={arch}") + for k in sorted(quant): + print(f" {prefix}{k} = {quant[k]!r}") + if quant.get("method") != "foldquant": + print("VIOLATION: quant.method must be 'foldquant'") + return 1 + + problems = 0 + sites = [] + for name, t in tensors.items(): + if t.tensor_type != gguf.GGMLQuantizationType.I8 or not name.endswith(".weight"): + continue + base = name[: -len(".weight")] + mod = "llm" if base.startswith("vlm.") else "action" + wbits = int(quant.get(f"{mod}_weight_bits", 8)) + shape = [int(s) for s in t.shape] # ggml ne order: [K_pack, N] + if len(shape) != 2: + print(f"VIOLATION: {name} is {len(shape)}-D"); problems += 1; continue + kpack, n_out = shape + k_in = kpack * 2 if wbits == 4 else kpack + ws = tensors.get(base + ".wscale") + asc = tensors.get(base + ".ascale") + ok = True + if ws is None or ws.tensor_type != gguf.GGMLQuantizationType.F32 or [int(s) for s in ws.shape] != [n_out]: + print(f"VIOLATION: {base}.wscale missing or not F32[{n_out}]"); ok = False + if asc is not None and (asc.tensor_type != gguf.GGMLQuantizationType.F32 or [int(s) for s in asc.shape] != [k_in]): + print(f"VIOLATION: {base}.ascale is not F32[{k_in}]"); ok = False + if k_in % 64 or n_out % 64: + print(f"VIOLATION: {name} K={k_in} N={n_out} must be multiples of 64"); ok = False + problems += not ok + bs = fq.rotation_block_for(k_in, int(quant.get(f"{mod}_rot_block_size", 64))) + sites.append((base, mod, k_in, n_out, wbits, asc is not None, bs)) + if not args.quiet: + print(f" site {base:40s} W{wbits} K={k_in:<5d} N={n_out:<5d} rot={bs:<3d} ascale={'yes' if asc is not None else 'no'}") + + # Fused groups share one activation transform. Which projections are fused + # follows the input they read: k and v always read the same tensor, and q + # joins them only on a self-attention site. The file carries no self/cross + # flag, so the split is inferred from K: on a cross-attention block q reads + # the hidden state and k/v the encoder, and their K differ (GR00T: 1536 vs + # 2048); when q's K equals k's, all three must agree, as vla.cpp's loader + # (fq_declare_fused) then fuses them. + groups = defaultdict(dict) + for base, _mod, k_in, *_ in sites: + m = GROUP.match(base) + if m: + groups[m.group(1)][m.group(2)] = (k_in, tensors.get(base + ".ascale")) + for g, members in groups.items(): + have = [v is not None for _, v in members.values()] + if any(have) and not all(have): + print(f"VIOLATION: {g}: .ascale present on some of q/k/v but not all"); problems += 1 + continue + if not all(have) or len(members) < 2: + continue + kq = members["q"][0] if "q" in members else None + fused = [n for n in ("q", "k", "v") if n in members and (n != "q" or kq == members.get("k", (kq,))[0])] + ref = np.asarray(members[fused[0]][1].data, dtype=np.float32) + for n in fused[1:]: + if not np.array_equal(ref, np.asarray(members[n][1].data, dtype=np.float32)): + print(f"VIOLATION: {g}: fused {'/'.join(fused)} .ascale vectors differ"); problems += 1; break + + n_llm = sum(1 for s in sites if s[1] == "llm") + n_act = len(sites) - n_llm + print(f"summary: {len(sites)} FoldQuant sites ({n_llm} llm, {n_act} action), {problems} violation(s)") + return 1 if problems else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/patch_ggml_openvino.py b/scripts/patch_ggml_openvino.py index 5fa4cc1..07e2d65 100755 --- a/scripts/patch_ggml_openvino.py +++ b/scripts/patch_ggml_openvino.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Thirteen fixes to the fetched ggml OpenVINO backend. +"""Thirteen fixes to the fetched ggml OpenVINO backend, and the FoldQuant hook. ggml-openvino is written against llama.cpp's graphs: one decoder-only transformer, one position input, an F16 KV cache. vla.cpp drives it with vision @@ -190,6 +190,14 @@ and a late grasp on a robot. GGML_OPENVINO_GPU_PRECISION=f32 puts it at 6.5e-5. Exposed rather than forced: f32 costs about 3x on this plugin, so src/backend.h defaults it for pi0 alone and an explicit setting still wins. + 14. op_table.{h,cpp}, ggml-openvino.cpp, ggml-quants.cpp, ggml-decoder.cpp - + the hook for FoldQuant. A FoldQuant GGUF builds two GGML_OP_CUSTOM nodes + per INT8/INT4 site, which ggml leaves to the CPU backend. The translator + that maps them onto OpenVINO ops lives in-tree (src/openvino/ + foldquant_ov.cpp, compiled into this target by vla.cpp's CMakeLists); this + hunk only registers it for GGML_OP_CUSTOM, lets supports_op accept + FoldQuant's nodes (and no other custom op) ahead of the type checks that + would refuse their INT8 tensors, and keeps an INT8 weight as an i8 constant. Idempotent - re-running on a patched tree is a no-op, so a reconfigure that re-populates the FetchContent source dir is safe either way. @@ -618,6 +626,56 @@ // defaults to erf, so the tanh variant must set its mode explicitly. {"GGML_UNARY_OP_GELU_ERF", op::translate_1to1_match_1_input },""", ), + ( + """ {"GGML_OP_ADD", op::translate_add },""", + """ {"GGML_OP_ADD", op::translate_add }, + // vla.cpp: FoldQuant's two custom nodes (src/openvino/foldquant_ov.cpp). + {"GGML_OP_CUSTOM", op::translate_vla_foldquant },""", + ), + ], + "ggml/src/ggml-openvino/openvino/op_table.h": [ + ( + """GGML_OP_CONVERTER(translate_add);""", + """GGML_OP_CONVERTER(translate_add); +GGML_OP_CONVERTER(translate_vla_foldquant); // vla.cpp: src/openvino/foldquant_ov.cpp""", + ), + ], + "ggml/src/ggml-openvino/ggml-openvino.cpp": [ + ( + """static ggml_openvino_op_support ggml_backend_openvino_device_supports_op_impl(ggml_backend_dev_t dev, const ggml_tensor * op) { + GGML_ASSERT(dev->reg != nullptr); +""", + """// vla.cpp: defined in src/openvino/foldquant_ov.cpp. +bool vla_foldquant_ov_supports(const ggml_tensor * op); + +static ggml_openvino_op_support ggml_backend_openvino_device_supports_op_impl(ggml_backend_dev_t dev, const ggml_tensor * op) { + GGML_ASSERT(dev->reg != nullptr); + + // vla.cpp: FoldQuant's custom nodes carry INT8 weights and an INT8 blob, + // which the type checks below would refuse; they have their own test. + if (op->op == GGML_OP_CUSTOM) { + if (vla_foldquant_ov_supports(op)) { + return {true, ""}; + } + return {false, "GGML_OP_CUSTOM other than FoldQuant's"}; + } +""", + ), + ], + "ggml/src/ggml-openvino/ggml-quants.cpp": [ + ( + """ if (tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_BF16) { + ov::element::Type element_type; + switch (tensor->type) {""", + """ // vla.cpp: GGML_TYPE_I8 is a FoldQuant site's codes, kept as an i8 constant. + if (tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_BF16 || + tensor->type == GGML_TYPE_I8) { + ov::element::Type element_type; + switch (tensor->type) { + case GGML_TYPE_I8: + element_type = ov::element::i8; + break;""", + ), ], "ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp": [ ( @@ -744,6 +802,11 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {""", ), + ( + """ static const std::set weight_types = {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0,""", + """ // vla.cpp: GGML_TYPE_I8 is a FoldQuant site's codes (see ggml-quants.cpp). + static const std::set weight_types = {GGML_TYPE_I8, GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0,""", + ), ], "ggml/src/ggml-openvino/utils.h": [ ( diff --git a/scripts/patch_ggml_sycl_ext_hook.py b/scripts/patch_ggml_sycl_ext_hook.py new file mode 100644 index 0000000..f9db29c --- /dev/null +++ b/scripts/patch_ggml_sycl_ext_hook.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Add one extension hook to the fetched ggml SYCL backend. + +The FoldQuant kernels for Intel GPUs live in src/sycl/vla_sycl_foldquant.cpp as +ordinary in-tree code; this is the ggml side, the SYCL counterpart of +scripts/patch_ggml_cuda_ext_hook.py. A FoldQuant GGUF builds two +GGML_OP_CUSTOM nodes per INT8/INT4 site, and ggml runs GGML_OP_CUSTOM on the +CPU backend only, so the SYCL backend has to offer one place where an external +implementation gets first refusal. + +What it changes (ggml/src/ggml-sycl/ggml-sycl.cpp only) +-------------------------------------------------------- + 1. Two exported function pointers, null by default: the compute hook and a + supports_op answer for the nodes it claims. + 2. One call to the first at the top of ggml_sycl_compute_forward. Returning + false means "not mine", and ggml runs the op exactly as before. + 3. One call to the second at the top of the supports_op switch, for + GGML_OP_CUSTOM only. + +With the pointers left null this is a no-op, so a hooked ggml behaves +identically to a stock one. + +Usage: scripts/patch_ggml_sycl_ext_hook.py [] +""" + +import pathlib +import sys + +MARKER = "vla.cpp: SYCL extension hook" + +HOOK_DECL = ( + """static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct ggml_tensor * dst) try { + GGML_SYCL_DEBUG("[SYCL] ggml_sycl_compute_forward: dst=%s, op=%s\\n", dst->name, ggml_op_name(dst->op)); + if (!g_sycl_loaded) return false; +""", + """// vla.cpp: SYCL extension hook. Null unless vla::sycl_register_foldquant_ops() +// ran; see src/sycl/vla_sycl_foldquant.cpp, which holds every kernel behind it. +extern "C" { +typedef bool (*ggml_sycl_ext_forward_t)(struct ggml_tensor * dst, void * queue); +__attribute__((visibility("default"))) ggml_sycl_ext_forward_t ggml_sycl_ext_forward = nullptr; +typedef bool (*ggml_sycl_ext_supports_t)(const struct ggml_tensor * op); +__attribute__((visibility("default"))) ggml_sycl_ext_supports_t ggml_sycl_ext_supports = nullptr; +} + +static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct ggml_tensor * dst) try { + GGML_SYCL_DEBUG("[SYCL] ggml_sycl_compute_forward: dst=%s, op=%s\\n", dst->name, ggml_op_name(dst->op)); + if (!g_sycl_loaded) return false; + if (ggml_sycl_ext_forward && ggml_sycl_ext_forward(dst, (void *) ctx.stream())) { + return true; + } +""", +) + +SUPPORTS = ( + """static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { +""", + """static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + // vla.cpp: SYCL extension hook - the custom nodes the extension claims. + if (op->op == GGML_OP_CUSTOM) { + return ggml_sycl_ext_supports && ggml_sycl_ext_supports(op); + } +""", +) + + +def main(): + src = pathlib.Path(sys.argv[1] if len(sys.argv) > 1 else ".").resolve() + path = src / "ggml/src/ggml-sycl/ggml-sycl.cpp" + if not path.exists(): + raise SystemExit(f"not a llama.cpp source tree: {src}") + + text = path.read_text() + if MARKER in text: + return # idempotent: re-configure over an already-patched tree + + for old, new in (HOOK_DECL, SUPPORTS): + n = text.count(old) + if n != 1: + raise SystemExit( + f"{path}: anchor found {n} times, expected 1. The pinned llama.cpp " + f"probably moved; re-check this anchor against the new tag.\n" + f"---\n{old[:400]}\n---" + ) + text = text.replace(old, new) + + path.write_text(text) + + +if __name__ == "__main__": + main() diff --git a/scripts/quantize_gguf.py b/scripts/quantize_gguf.py index 130ea19..f364d92 100644 --- a/scripts/quantize_gguf.py +++ b/scripts/quantize_gguf.py @@ -21,6 +21,13 @@ not benefit and can lose accuracy). python scripts/quantize_gguf.py --in model-bf16.gguf --out model-q8_0.gguf --type Q8_0 + +This is the stock ggml repack: block-32 absmax weights, float activations, +dequantized inside ggml_mul_mat, runs on every backend. It is not FoldQuant. +A FoldQuant GGUF (docs/QUANTIZATION.md: INT8/INT4 codes + sidecar scales executed +by the in-tree integer kernels on CUDA/CPU, read back as float weights elsewhere) is produced from a FoldQuantVLA +quantized model by scripts/convert_quantized_model_to_gguf.py or, uncalibrated, +by scripts/foldquant_fake_export.py; this script refuses such a file. """ import argparse @@ -93,6 +100,9 @@ def main() -> None: r = gguf.GGUFReader(args.src) arch = r.fields["general.architecture"].contents() + if f"{arch}.quant.method" in r.fields: + raise SystemExit(f"{args.src} is a FoldQuant GGUF; its INT8 sites are not ggml block types " + "and cannot be repacked (see docs/QUANTIZATION.md)") w = gguf.GGUFWriter(args.dst, arch) copy_kv(r, w) diff --git a/src/backend.h b/src/backend.h index 7f61a95..c786426 100644 --- a/src/backend.h +++ b/src/backend.h @@ -64,18 +64,17 @@ #include #include #endif +#include // setenv / _putenv_s #if defined(GGML_USE_SYCL) || defined(GGML_USE_OPENVINO) || \ defined(GGML_USE_HEXAGON) || defined(GGML_USE_OPENCL) -#include // setenv / _putenv_s #include #endif namespace vla { -#if defined(GGML_USE_SYCL) || defined(GGML_USE_OPENVINO) || \ - defined(GGML_USE_HEXAGON) || defined(GGML_USE_OPENCL) // setenv is POSIX. _putenv_s has no "do not overwrite" mode, so check first. // Empty counts as unset; an empty KEY= in a compose file is not a choice. +// Used by the SYCL/OpenVINO/Hexagon/OpenCL ladders and by the FoldQuant reference switch. inline void setenv_default(const char * key, const char * val) { #ifdef _WIN32 size_t len = 0; @@ -89,7 +88,6 @@ inline void setenv_default(const char * key, const char * val) { setenv(key, val, /*overwrite=*/1); #endif } -#endif /// Outcome of @ref backend_init. @c handle is null only if even the CPU backend /// failed to come up, which callers treat as a fatal load error. diff --git a/src/cuda/vla_cuda_bf16.cu b/src/cuda/vla_cuda_bf16.cu index 70b3766..64174ab 100644 --- a/src/cuda/vla_cuda_bf16.cu +++ b/src/cuda/vla_cuda_bf16.cu @@ -39,6 +39,7 @@ #include "ggml.h" #include "env_flag.h" +#include "cuda/vla_cuda_ext.h" #include #include @@ -740,7 +741,7 @@ namespace vla { // Called once, after the CUDA backend is up. Idempotent. void cuda_register_bf16_ops() { - ggml_cuda_ext_forward = vla_cuda_bf16_forward; + cuda_ext_add_handler(vla_cuda_bf16_forward); // shares the hook with the FoldQuant handler // Fusion runs in ggml_backend_cuda_graph_compute, upstream of // ggml_cuda_compute_forward, so the pointer above never sees a fused node. diff --git a/src/cuda/vla_cuda_ext.cu b/src/cuda/vla_cuda_ext.cu new file mode 100644 index 0000000..03f0f71 --- /dev/null +++ b/src/cuda/vla_cuda_ext.cu @@ -0,0 +1,59 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// The ggml CUDA extension hook is one function pointer. Several in-tree kernel +// families ride it (BF16 activations, FoldQuant linears), so this dispatcher is +// the only thing ever stored there: each family registers a handler that +// returns false, with no side effects, for nodes that are not its own. + +#include "ggml.h" + +#include + +extern "C" { +typedef bool (*ggml_cuda_ext_forward_t)(struct ggml_tensor * dst, void * stream); +extern ggml_cuda_ext_forward_t ggml_cuda_ext_forward; +} + +namespace vla { + +using cuda_ext_fn = bool (*)(ggml_tensor *, void *); + +namespace { + +constexpr int MAX_HANDLERS = 4; +cuda_ext_fn g_handlers[MAX_HANDLERS] = {}; +int g_n_handlers = 0; + +extern "C" bool vla_cuda_ext_dispatch(ggml_tensor * dst, void * stream) { + for (int i = 0; i < g_n_handlers; ++i) + if (g_handlers[i](dst, stream)) + return true; + return false; +} + +} // namespace + +// Idempotent per handler. Order of registration is irrelevant: handlers +// self-select on the node (type, op, magic word) and never overlap. +void cuda_ext_add_handler(cuda_ext_fn fn) { + for (int i = 0; i < g_n_handlers; ++i) + if (g_handlers[i] == fn) + return; + if (g_n_handlers < MAX_HANDLERS) + g_handlers[g_n_handlers++] = fn; + ggml_cuda_ext_forward = vla_cuda_ext_dispatch; +} + +} // namespace vla diff --git a/src/cuda/vla_cuda_ext.h b/src/cuda/vla_cuda_ext.h new file mode 100644 index 0000000..6a3561e --- /dev/null +++ b/src/cuda/vla_cuda_ext.h @@ -0,0 +1,27 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "ggml.h" + +namespace vla { + +using cuda_ext_fn = bool (*)(ggml_tensor *, void *); + +// Install a first-refusal handler for ggml_cuda_compute_forward (see +// vla_cuda_ext.cu). Idempotent. +void cuda_ext_add_handler(cuda_ext_fn fn); + +} // namespace vla diff --git a/src/cuda/vla_cuda_foldquant.cu b/src/cuda/vla_cuda_foldquant.cu new file mode 100644 index 0000000..cc74a46 --- /dev/null +++ b/src/cuda/vla_cuda_foldquant.cu @@ -0,0 +1,312 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant nodes on the CUDA backend. A site is two GGML_OP_CUSTOM nodes +// (layers/fq_linear.h) whose userdata points at an FqActSpec / FqGemmSpec; the +// magic word in that struct is how this handler recognises its nodes. +// +// Everything is validated before anything is launched. A node that fails the +// contract is declined (return false), which makes ggml abort loudly on the +// unsupported op instead of silently computing something else. +// +// VLA_FQ_CPU_REF=1 routes every node through the CPU reference on host copies +// of the device tensors: slow, but a byte-exact A/B against the kernels. + +#include "foldquant.h" +#include +#include +#include "foldquant_ref.h" +#include "env_flag.h" +#include "kernels/foldquant/fq_kernels.h" +#include "cuda/vla_cuda_ext.h" + +#include "ggml.h" + +#include + +#include +#include +#include + +namespace { + +// Layout of the op_params ggml_custom_4d writes (ggml-impl.h, private). Pinned +// by tests/test_foldquant_cuda_op.cpp, which checks that the CPU op and this +// decoder see the same userdata. +struct custom_op_params_mirror { + void * fun; + int n_tasks; + void * userdata; +}; + +const void * node_userdata(const ggml_tensor * dst) { + if (dst->op != GGML_OP_CUSTOM) return nullptr; + custom_op_params_mirror p; + std::memcpy(&p, dst->op_params, sizeof(p)); + return p.userdata; +} + +uint32_t magic_of(const void * ud) { + uint32_t m = 0; + if (ud) std::memcpy(&m, ud, sizeof(m)); + return m; +} + +bool trace() { + static const bool t = vla::env_flag("VLA_FQ_TRACE"); + return t; +} + +bool cpu_ref() { + static const bool r = vla::env_flag("VLA_FQ_CPU_REF"); + return r; +} + +// --- CPU reference on host copies ------------------------------------------ + +struct HostCopy { + ggml_tensor t; + std::vector bytes; +}; + +bool host_copy(const ggml_tensor * src, HostCopy & out, cudaStream_t stream) { + out.t = *src; + out.bytes.resize(ggml_nbytes(src)); + if (cudaMemcpyAsync(out.bytes.data(), src->data, out.bytes.size(), cudaMemcpyDeviceToHost, stream) != cudaSuccess) + return false; + out.t.data = out.bytes.data(); + return true; +} + +template +bool run_on_host(ggml_tensor * dst, int n_src, Fn fn, cudaStream_t stream) { + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + if (cudaStreamIsCapturing(stream, &cap) == cudaSuccess && cap != cudaStreamCaptureStatusNone) { + std::fprintf(stderr, "vla(fq): VLA_FQ_CPU_REF cannot run inside CUDA-graph capture; " + "set GGML_CUDA_DISABLE_GRAPHS=1 (foldquant_check_backend does so at load)\n"); + return false; + } + HostCopy srcs[GGML_MAX_SRC]; + ggml_tensor d = *dst; + for (int i = 0; i < n_src; ++i) { + if (!dst->src[i]) { d.src[i] = nullptr; continue; } + if (!host_copy(dst->src[i], srcs[i], stream)) return false; + d.src[i] = &srcs[i].t; + } + std::vector out(ggml_nbytes(dst)); + d.data = out.data(); + if (cudaStreamSynchronize(stream) != cudaSuccess) return false; + fn(&d); + return cudaMemcpyAsync(dst->data, out.data(), out.size(), cudaMemcpyHostToDevice, stream) == cudaSuccess; +} + +// VLA_FQ_CHECK=1: after each kernel, recompute the node with the CPU reference +// on host copies of the same inputs and report any byte difference. Slow; +// a diagnostic for a mismatch that the unit tests' shapes do not reproduce. +bool check_mode() { + static const bool c = vla::env_flag("VLA_FQ_CHECK"); + return c; +} + +template +void check_against_host(ggml_tensor * dst, int n_src, Fn fn, cudaStream_t stream, size_t cmp_bytes_per_row, size_t row_bytes) { + HostCopy srcs[GGML_MAX_SRC]; + ggml_tensor d = *dst; + for (int i = 0; i < n_src; ++i) { + if (!dst->src[i]) { d.src[i] = nullptr; continue; } + if (!host_copy(dst->src[i], srcs[i], stream)) return; + d.src[i] = &srcs[i].t; + } + std::vector got(ggml_nbytes(dst)), want(ggml_nbytes(dst)); + if (cudaMemcpyAsync(got.data(), dst->data, got.size(), cudaMemcpyDeviceToHost, stream) != cudaSuccess) return; + if (cudaStreamSynchronize(stream) != cudaSuccess) return; + d.data = want.data(); + fn(&d); + const int64_t rows = dst->ne[1]; + int64_t bad_rows = 0, first = -1; + size_t first_off = 0; + for (int64_t r = 0; r < rows; ++r) { + const uint8_t * a = got.data() + (size_t) r * row_bytes; + const uint8_t * b = want.data() + (size_t) r * row_bytes; + for (size_t k = 0; k < cmp_bytes_per_row; ++k) + if (a[k] != b[k]) { ++bad_rows; if (first < 0) { first = r; first_off = k; } break; } + } + if (bad_rows) { + std::printf("vla(fq) CHECK %-40s ne=[%lld,%lld] MISMATCH rows %lld/%lld (first row %lld byte %zu)\n", + ggml_get_name(dst), (long long) dst->ne[0], (long long) dst->ne[1], (long long) bad_rows, + (long long) rows, (long long) first, first_off); + if (dst->type == GGML_TYPE_F32) { + const float * a = (const float *) (got.data() + (size_t) first * row_bytes); + const float * b = (const float *) (want.data() + (size_t) first * row_bytes); + std::printf(" kernel:"); for (int k = 0; k < 6; ++k) std::printf(" %.7g", a[k]); std::printf("\n"); + std::printf(" host :"); for (int k = 0; k < 6; ++k) std::printf(" %.7g", b[k]); std::printf("\n"); + } + } + else + std::printf("vla(fq) CHECK %-40s ne=[%lld,%lld] ok\n", ggml_get_name(dst), (long long) dst->ne[0], (long long) dst->ne[1]); +} + +// Rows contiguous and 16-byte aligned, any row stride, higher dims packed: what +// the prologue reads without a ggml_cont in front of it. +// VLA_FQ_PREFETCH_MB: how much of the next site's weights a GEMM prefetches +// into L2. Off by default: on Orin (4 MB L2) a 2 MB prefetch made the GEMMs +// 12% slower in the model (extra DRAM traffic, no hit-rate gain). +int64_t prefetch_bytes() { + static const int64_t b = [] { + const char * e = std::getenv("VLA_FQ_PREFETCH_MB"); + return (e && *e) ? (int64_t) (std::atof(e) * 1024.0 * 1024.0) : (int64_t) 0; + }(); + return b; +} + +bool contiguous_f32_rows(const ggml_tensor * t, int64_t K) { + return t && t->type == GGML_TYPE_F32 && t->ne[0] == K && t->nb[0] == sizeof(float) && + t->nb[1] % 16 == 0 && t->nb[2] == t->nb[1] * (size_t) t->ne[1] && t->nb[3] == t->nb[2] * (size_t) t->ne[2]; +} + +// --- the two nodes ---------------------------------------------------------- + +bool forward_act(ggml_tensor * dst, const vla::FqActSpec & s, cudaStream_t stream) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * g = nullptr, * as = nullptr; + vla::fq_act_srcs(dst, s, &g, &as); + const int64_t K = s.K; + const int64_t rb = vla::fq_act_row_bytes(K, s.abits); + if (!contiguous_f32_rows(x, K) || dst->type != GGML_TYPE_I8 || dst->ne[0] != rb || + (as && !(as->type == GGML_TYPE_F32 && as->ne[0] == K)) || + (g && !(g->type == GGML_TYPE_F32 && g->ne[0] == K)) || + (s.has_gamma && !g) || (s.has_ascale && !as)) { + std::fprintf(stderr, "vla(fq): act node %s violates the contract (x %s ne0=%lld, dst ne0=%lld)\n", + ggml_get_name(dst), x ? ggml_type_name(x->type) : "null", + x ? (long long) x->ne[0] : 0ll, (long long) dst->ne[0]); + return false; + } + const int64_t M = ggml_nelements(x) / K; + if (trace()) + std::printf("vla(fq): act %s M=%lld K=%lld A%d rot%d%s%s\n", ggml_get_name(dst), (long long) M, + (long long) K, s.abits, s.rot_block, g ? " +rmsnorm" : "", as ? " +ascale" : ""); + + if (cpu_ref()) + return run_on_host(dst, 3, [&](ggml_tensor * d) { vla::fq_act_cpu(d, 0, 1, (void *) &s); }, stream); + + vla::fq::ActArgs a; + a.x = (const float *) x->data; + a.x_stride = (int64_t) (x->nb[1] / sizeof(float)); + a.ascale = as ? (const float *) as->data : nullptr; + a.gamma = (s.has_gamma && g) ? (const float *) g->data : nullptr; + a.blob = (int8_t *) dst->data; + a.row_bytes = rb; + a.M = M; a.K = K; + a.abits = s.abits; a.rot_block = s.rot_block; a.fold_before = s.fold_before; + a.clip = s.clip; a.eps = s.eps; + a.inv_sqrt_bs = vla::fqref::inv_sqrt_block(s.rot_block > 1 ? s.rot_block : 1); + const cudaError_t e = vla::fq::launch_act(a, stream); + if (e == cudaErrorNotSupported) + return run_on_host(dst, 3, [&](ggml_tensor * d) { vla::fq_act_cpu(d, 0, 1, (void *) &s); }, stream); + if (e != cudaSuccess) { + std::fprintf(stderr, "vla(fq): act launch failed: %s\n", cudaGetErrorString(e)); + return false; + } + if (check_mode()) + check_against_host(dst, 3, [&](ggml_tensor * d) { vla::fq_act_cpu(d, 0, 1, (void *) &s); }, stream, + (size_t) vla::fq_act_kpack(K, s.abits) + 4, (size_t) rb); + return true; +} + +bool forward_gemm(ggml_tensor * dst, const vla::FqGemmSpec & s, cudaStream_t stream) { + const ggml_tensor * w = dst->src[0]; + const ggml_tensor * xq = dst->src[1]; + const ggml_tensor * ws = dst->src[2]; + const ggml_tensor * b = dst->src[3]; + const ggml_tensor * r = dst->src[4]; + const int64_t K = s.K, N = s.N; + const int64_t kpw = vla::fq_w_kpack(K, s.wbits); + if (!w || w->type != GGML_TYPE_I8 || w->ne[0] != kpw || w->ne[1] != N || !ggml_is_contiguous(w) || + !xq || xq->type != GGML_TYPE_I8 || !ggml_is_contiguous(xq) || + !ws || ws->type != GGML_TYPE_F32 || ws->ne[0] != N || + (b && !(b->type == GGML_TYPE_F32 && b->ne[0] == N)) || + (r && !(r->type == GGML_TYPE_F32 && ggml_is_contiguous(r) && ggml_nelements(r) == ggml_nelements(dst) && r->ne[0] == N)) || + dst->type != GGML_TYPE_F32 || dst->ne[0] != N || !ggml_is_contiguous(dst)) { + std::fprintf(stderr, "vla(fq): gemm node %s violates the contract\n", ggml_get_name(dst)); + return false; + } + const int abits = (xq->ne[0] == K + vla::FQ_ACT_TAIL) ? 8 : 4; + if (xq->ne[0] != vla::fq_act_row_bytes(K, abits)) { + std::fprintf(stderr, "vla(fq): gemm node %s: blob row of %lld bytes does not match K=%lld\n", + ggml_get_name(dst), (long long) xq->ne[0], (long long) K); + return false; + } + const int64_t M = xq->ne[1]; + if (trace()) + std::printf("vla(fq): gemm %s M=%lld N=%lld K=%lld W%dA%d%s\n", ggml_get_name(dst), (long long) M, + (long long) N, (long long) K, s.wbits, abits, b ? " +bias" : ""); + + if (cpu_ref()) + return run_on_host(dst, 5, [&](ggml_tensor * d) { vla::fq_gemm_cpu(d, 0, 1, (void *) &s); }, stream); + + vla::fq::GemmArgs g; + g.w = (const int8_t *) w->data; + g.blob = (const int8_t *) xq->data; + g.wscale = (const float *) ws->data; + g.bias = b ? (const float *) b->data : nullptr; + g.res = r ? (const float *) r->data : nullptr; + if (s.next_w && s.next_w->data && prefetch_bytes() > 0) { + g.pf = (const int8_t *) s.next_w->data; + g.pf_bytes = std::min((int64_t) ggml_nbytes(s.next_w), prefetch_bytes()); + } + g.y = (float *) dst->data; + g.head_dim = s.head_dim; g.heads = s.heads; g.vmask = s.vmask; + g.M = M; g.N = N; g.K = K; + g.row_bytes = xq->ne[0]; + g.wbits = s.wbits; g.abits = abits; + const cudaError_t e = vla::fq::launch_gemm(g, stream); + if (e == cudaErrorNotSupported) { + static bool warned = false; + if (!warned) { + std::printf("vla(fq): W%dA%d GEMM has no CUDA kernel yet; running the CPU reference for those sites\n", + s.wbits, abits); + warned = true; + } + return run_on_host(dst, 5, [&](ggml_tensor * d) { vla::fq_gemm_cpu(d, 0, 1, (void *) &s); }, stream); + } + if (e != cudaSuccess) { + std::fprintf(stderr, "vla(fq): gemm launch failed: %s\n", cudaGetErrorString(e)); + return false; + } + if (check_mode()) + check_against_host(dst, 5, [&](ggml_tensor * d) { vla::fq_gemm_cpu(d, 0, 1, (void *) &s); }, stream, + (size_t) N * sizeof(float), (size_t) N * sizeof(float)); + return true; +} + +} // namespace + +extern "C" bool vla_cuda_foldquant_forward(ggml_tensor * dst, void * stream_v) { + if (!dst || dst->op != GGML_OP_CUSTOM) + return false; + const void * ud = node_userdata(dst); + switch (magic_of(ud)) { + case vla::FQ_ACT_MAGIC: return forward_act (dst, *(const vla::FqActSpec *) ud, (cudaStream_t) stream_v); + case vla::FQ_GEMM_MAGIC: return forward_gemm(dst, *(const vla::FqGemmSpec *) ud, (cudaStream_t) stream_v); + default: return false; + } +} + +namespace vla { + +void cuda_register_foldquant_ops() { + cuda_ext_add_handler(vla_cuda_foldquant_forward); +} + +} // namespace vla diff --git a/src/cuda/vla_cuda_ops.h b/src/cuda/vla_cuda_ops.h index b6085b6..b16c3b3 100644 --- a/src/cuda/vla_cuda_ops.h +++ b/src/cuda/vla_cuda_ops.h @@ -14,17 +14,22 @@ #pragma once -// Registration for the in-tree BF16 CUDA kernels (src/cuda/vla_cuda_bf16.cu). +// Registration for the in-tree CUDA kernel families that ride the ggml +// extension hook: the BF16 activation ops (src/cuda/vla_cuda_bf16.cu) and the +// FoldQuant INT8/INT4 linears (src/cuda/vla_cuda_foldquant.cu). Both install +// through one dispatcher (src/cuda/vla_cuda_ext.cu) so they compose. // -// Off every other build: without CUDA there is no hook to install and the BF16 -// activation path is unreachable anyway, so this compiles to nothing. +// Off every other build: without CUDA there is no hook to install and these +// paths are unreachable anyway, so this compiles to nothing. namespace vla { #ifdef GGML_USE_CUDA void cuda_register_bf16_ops(); +void cuda_register_foldquant_ops(); #else inline void cuda_register_bf16_ops() {} +inline void cuda_register_foldquant_ops() {} #endif } // namespace vla diff --git a/src/foldquant.cpp b/src/foldquant.cpp new file mode 100644 index 0000000..4b7c9df --- /dev/null +++ b/src/foldquant.cpp @@ -0,0 +1,444 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "foldquant.h" + +#include "backend.h" +#include "env_flag.h" +#include "gguf_reader.h" +#include "loader.h" +#include "cuda/vla_cuda_ops.h" +#include "foldquant_ref.h" +#ifdef VLA_FQ_SYCL +#include "sycl/vla_sycl_foldquant.h" +#endif + +#include +#include +#include +#include +#include +#include + +namespace vla { + +namespace { + +constexpr size_t NAME_CAP = 256; + +std::string key(const char * prefix, const char * k) { + return std::string(prefix) + ".quant." + k; +} + +// "o:8,down:8" -> {o: 8, down: 8} +std::map parse_site_bits(const std::string & s) { + std::map out; + std::stringstream ss(s); + std::string item; + while (std::getline(ss, item, ',')) { + const size_t c = item.find(':'); + if (c == std::string::npos || c == 0) continue; + out[item.substr(0, c)] = std::atoi(item.c_str() + c + 1); + } + return out; +} + +void parse_module(const gguf_reader & g, const char * prefix, const char * mod, FqModuleSpec & m) { + const std::string p = std::string(prefix) + ".quant." + mod + "_"; + auto u32 = [&](const char * k, int def) { + const std::string kk = p + k; + return g.has(kk.c_str()) ? (int) g.u32(kk.c_str()) : def; + }; + m.wbits = u32("weight_bits", m.wbits); + m.abits = u32("act_bits", m.abits); + m.rot_block = u32("rot_block_size", m.rot_block); + m.scheme = g.str(key(prefix, (std::string("scheme_") + mod).c_str()).c_str()); +} + +} // namespace + +int FqModuleSpec::wbits_for(const char * site_key) const { + const auto it = site_bits.find(site_key); + return it != site_bits.end() ? it->second : wbits; +} + +int FqModuleSpec::abits_for(const char * site_key) const { + const auto it = site_bits.find(site_key); + return it != site_bits.end() ? it->second : abits; +} + +int fq_rot_block_for(int64_t K, int nominal) { + int bs = nominal; + while (bs > 1 && (K % bs) != 0) + bs >>= 1; + return bs < 2 ? 1 : bs; +} + +bool foldquant_present(const gguf_reader & g, const char * prefix) { + const std::string k = key(prefix, "method"); + return g.has(k.c_str()) && g.str(k.c_str()) == "foldquant"; +} + +FoldQuantSpec foldquant_parse(const gguf_reader & g, const char * prefix) { + FoldQuantSpec fq; + if (!foldquant_present(g, prefix)) + return fq; + fq.present = true; + fq.method = g.str(key(prefix, "method").c_str()); + fq.applied_at = g.str(key(prefix, "applied_at").c_str()); + fq.provenance = g.str(key(prefix, "provenance").c_str()); + + parse_module(g, prefix, "llm", fq.llm); + parse_module(g, prefix, "action", fq.action); + + const std::string fo = g.str(key(prefix, "action_fold_order").c_str()); + fq.action.fold_before = (fo == "before"); + fq.llm.fold_before = false; // the LLM never ships an ascale + + const std::string clipk = key(prefix, "act_clip_ratio"); + if (g.has(clipk.c_str())) { + const float c = g.f32(clipk.c_str()); + if (c > 0.0f) fq.llm.clip = c; + } + fq.llm.site_bits = parse_site_bits(g.str(key(prefix, "site_bits").c_str())); + + std::printf("vla(%s): FoldQuant GGUF (%s): llm=%s W%dA%d rot%d%s action=%s W%dA%d rot%d fold=%s%s\n", + prefix, fq.applied_at.c_str(), + fq.llm.scheme.c_str(), fq.llm.wbits, fq.llm.abits, fq.llm.rot_block, + fq.llm.site_bits.empty() ? "" : " (site_bits)", + fq.action.scheme.c_str(), fq.action.wbits, fq.action.abits, fq.action.rot_block, + fq.action.fold_before ? "before" : "after", + fq.llm.clip != 1.0f ? " clip" : ""); + if (!fq.provenance.empty()) + std::printf("vla(%s): FoldQuant provenance: %s\n", prefix, fq.provenance.c_str()); + return fq; +} + +bool foldquant_check_backend(const char * tag, const Backend & b, FoldQuantSpec & fq, bool weight_dtype_set) { + if (!fq.present) + return true; + const char * name = b.handle ? ggml_backend_name(b.handle) : ""; + const bool is_cpu = std::strcmp(name, "CPU") == 0; +#ifdef VLA_FQ_SYCL + // SYCL: the integer kernels in src/sycl/vla_sycl_foldquant.cpp, behind the + // hook scripts/patch_ggml_sycl_ext_hook.py adds. They write plain [N][T] + // outputs, so the head-laid-out epilogue is off. + if (std::strncmp(name, "SYCL", 4) == 0 && !env_flag("VLA_FQ_DEQUANT")) { + fq.llm.no_heads = fq.action.no_heads = true; + sycl_register_foldquant_ops(); + if (weight_dtype_set) + std::printf("%s: --weight-dtype applies to the float tensors; FoldQuant sites stay INT%d/INT%d\n", + tag, fq.llm.wbits, fq.action.wbits); + std::printf("%s: FoldQuant SYCL kernels registered on '%s'\n", tag, name); + return true; + } +#endif + bool native_ov = false; +#ifdef VLA_FQ_OPENVINO + // ggml's OpenVINO backend runs the nodes through the in-tree translator on the + // CPU and GPU plugins; the NPU compiler takes no such graph, so it dequantizes. + if (std::strcmp(name, "OPENVINO") == 0) { + const char * dev = std::getenv("GGML_OPENVINO_DEVICE"); + native_ov = !(dev && std::strcmp(dev, "NPU") == 0); + } +#endif + if (native_ov && !env_flag("VLA_FQ_DEQUANT")) { + // The head-laid-out epilogue writes a layout the graph reads through views + // of raw memory, which a translated graph does not have. + fq.llm.no_heads = fq.action.no_heads = true; + std::printf("%s: FoldQuant on OpenVINO (%s): INT%d/INT%d weights, per-token INT%d/INT%d activations\n", + tag, std::getenv("GGML_OPENVINO_DEVICE") ? std::getenv("GGML_OPENVINO_DEVICE") : "CPU", + fq.llm.wbits, fq.action.wbits, fq.llm.abits, fq.action.abits); + return true; + } + if ((!b.is_cuda && !is_cpu) || env_flag("VLA_FQ_DEQUANT")) { + // No FoldQuant kernels on this backend (or asked for): every site is read + // back as a float weight in the resident GEMM type and runs the stock path. + fq.llm.dequant = fq.action.dequant = true; + std::printf("%s: FoldQuant sites read back as float weights on '%s' (%s): weights keep the " + "W%d/W%d rounding, activations stay float\n", + tag, name, (!b.is_cuda && !is_cpu) ? "no integer kernels on this backend" : "VLA_FQ_DEQUANT=1", + fq.llm.wbits, fq.action.wbits); + return true; + } + if (weight_dtype_set) + std::printf("%s: --weight-dtype applies to the float tensors; FoldQuant sites stay INT%d/INT%d\n", + tag, fq.llm.wbits, fq.action.wbits); + if (b.is_cuda) { + if (env_flag("VLA_FQ_CPU_REF")) { + // The staging shim synchronizes the stream, which CUDA-graph + // capture forbids; ggml reads this switch before its first compute. + setenv_default("GGML_CUDA_DISABLE_GRAPHS", "1"); + std::printf("%s: VLA_FQ_CPU_REF=1 - FoldQuant nodes run the CPU reference on host copies " + "(CUDA graphs disabled)\n", tag); + } + cuda_register_foldquant_ops(); + std::printf("%s: FoldQuant CUDA kernels registered\n", tag); + } else { + std::printf("%s: FoldQuant on the CPU reference path (exact, slow)\n", tag); + } + return true; +} + +void fq_dequant_rows(const int8_t * codes, const float * wscale, const float * ascale, + int64_t K, int64_t N, int wbits, int rot_block, bool fold_before, float * out) { + // The activation path is x' = R(x / a) (fold before) or R(x) / a (after), with R + // the block-normalised Sylvester-Hadamard butterfly, which is symmetric and + // orthonormal; so W x' = (W R diag(1/a)) x or (W diag(1/a) R) x, and row n of + // the float weight is R(w_n) / a or R(w_n / a). + const int64_t kpack = fq_w_kpack(K, wbits); + const float inv = rot_block > 1 ? fqref::inv_sqrt_block(rot_block) : 1.0f; + std::vector q((size_t) K); + for (int64_t n = 0; n < N; ++n) { + const int8_t * row = codes + (size_t) n * kpack; + if (wbits == 4) + fqref::unpack_nibbles(row, K, q.data()); + else + std::copy(row, row + K, q.begin()); + float * w = out + (size_t) n * K; + for (int64_t k = 0; k < K; ++k) + w[k] = (float) q[(size_t) k] * wscale[n]; + if (ascale && !fold_before) + for (int64_t k = 0; k < K; ++k) w[k] = w[k] / ascale[k]; + if (rot_block > 1) + fqref::fwht_row(w, K, rot_block, inv); + if (ascale && fold_before) + for (int64_t k = 0; k < K; ++k) w[k] = w[k] / ascale[k]; + } +} + +namespace { + +// Dequant mode: register `.weight` with the loader as a float [K, N] weight +// rebuilt from the codes, wscale and ascale. Declares nothing itself. +bool register_dequant(WeightLoader & L, const FqModuleSpec & mod, const char * site_key, const std::string & site) { + const std::string wname = site + ".weight"; + const ggml_tensor * meta = L.reader().meta(wname.c_str()); + const int wbits = mod.wbits_for(site_key); + if (wbits != 8 && wbits != 4) { + std::fprintf(stderr, "vla: %s: unsupported FoldQuant weight width %d\n", site.c_str(), wbits); + L.fail("FoldQuant widths"); + return false; + } + const int64_t K = wbits == 4 ? 2 * meta->ne[0] : meta->ne[0]; + const int64_t N = meta->ne[1]; + const std::string sname = site + ".wscale", aname = site + ".ascale"; + const ggml_tensor * sm = L.reader().meta(sname.c_str()); + const ggml_tensor * am = L.reader().meta(aname.c_str()); + if (ggml_n_dims(meta) != 2 || K % 64 != 0 || !sm || sm->ne[0] != N || (am && am->ne[0] != K)) { + std::fprintf(stderr, "vla: %s: malformed FoldQuant site (K=%lld N=%lld, wscale %s, ascale %s)\n", + site.c_str(), (long long) K, (long long) N, sm ? "ok" : "missing", am ? "present" : "absent"); + L.fail("FoldQuant site"); + return false; + } + const int rot = fq_rot_block_for(K, mod.rot_block); + const bool before = mod.fold_before; + gguf_reader & g = L.reader(); + const bool has_a = am != nullptr; + L.as_float(wname, K, N, [&g, wname, sname, aname, has_a, K, N, wbits, rot, before](std::vector & out) { + const std::vector codes = g.read_convert(wname.c_str(), GGML_TYPE_I8); + const std::vector ws = g.read_f32(sname.c_str()); + const std::vector as = has_a ? g.read_f32(aname.c_str()) : std::vector(); + if ((int64_t) codes.size() != fq_w_kpack(K, wbits) * N || (int64_t) ws.size() != N || + (has_a && (int64_t) as.size() != K)) + return false; + out.assign((size_t) (K * N), 0.0f); + fq_dequant_rows((const int8_t *) codes.data(), ws.data(), has_a ? as.data() : nullptr, + K, N, wbits, rot, before, out.data()); + return true; + }); + return true; +} + +bool fill_specs(WeightLoader & L, const FqModuleSpec & mod, const char * site_key, const char * site, + ggml_tensor * gamma, float eps, FqLinear & r) { + const int wbits = mod.wbits_for(site_key); + const int abits = mod.abits_for(site_key); + if ((wbits != 8 && wbits != 4) || (abits != 8 && abits != 4)) { + std::fprintf(stderr, "vla: %s: unsupported FoldQuant widths W%dA%d\n", site, wbits, abits); + L.fail("FoldQuant widths"); + return false; + } + const int64_t K = wbits == 4 ? 2 * r.w->ne[0] : r.w->ne[0]; + const int64_t N = r.w->ne[1]; + if (ggml_n_dims(r.w) != 2 || K % 64 != 0 || N % 64 != 0) { + std::fprintf(stderr, "vla: %s: FoldQuant site needs 2-D weight with K%%64==0 and N%%64==0, got K=%lld N=%lld\n", + site, (long long) K, (long long) N); + L.fail("FoldQuant shape"); + return false; + } + if (r.wscale->ne[0] != N || ggml_n_dims(r.wscale) != 1) { + std::fprintf(stderr, "vla: %s.wscale must be F32[%lld]\n", site, (long long) N); + L.fail("FoldQuant wscale"); + return false; + } + if (r.ascale && (r.ascale->ne[0] != K || ggml_n_dims(r.ascale) != 1)) { + std::fprintf(stderr, "vla: %s.ascale must be F32[%lld]\n", site, (long long) K); + L.fail("FoldQuant ascale"); + return false; + } + if (gamma && gamma->ne[0] != K) { + std::fprintf(stderr, "vla: %s: fused norm gamma has %lld entries, K=%lld\n", + site, (long long) gamma->ne[0], (long long) K); + L.fail("FoldQuant gamma"); + return false; + } + r.gamma = gamma; + r.act.K = K; + r.act.abits = abits; + r.act.rot_block = fq_rot_block_for(K, mod.rot_block); + if (r.act.rot_block > 64) { + std::fprintf(stderr, "vla: %s: rotation block %d exceeds the 64-element chunk the kernels rotate\n", + site, r.act.rot_block); + L.fail("FoldQuant rot_block"); + return false; + } + r.act.fold_before = mod.fold_before; + r.act.has_gamma = gamma != nullptr; + r.act.has_ascale = r.ascale != nullptr; + r.act.clip = abits == 4 ? mod.clip : 1.0f; + r.act.eps = eps; + r.gemm.K = K; + r.gemm.N = N; + r.gemm.wbits = wbits; + r.heads_off = mod.no_heads; + return true; +} + +} // namespace + +FqLinear fq_declare_linear(WeightLoader & L, const FqModuleSpec & mod, const char * site_key, + bool has_bias, ggml_tensor * gamma, float eps, const char * site_fmt, ...) { + char site[NAME_CAP]; + va_list ap; + va_start(ap, site_fmt); + const int n = std::vsnprintf(site, sizeof(site), site_fmt, ap); + va_end(ap); + FqLinear r; + if (n < 0 || (size_t) n >= sizeof(site)) { + L.fail("FoldQuant site name too long"); + return r; + } + + const std::string wname = std::string(site) + ".weight"; + const ggml_tensor * meta = L.reader().meta(wname.c_str()); + if (!meta || meta->type != GGML_TYPE_I8) + return r; // not a FoldQuant site: caller declares it as a float GEMM + if (mod.dequant) { + register_dequant(L, mod, site_key, site); + return r; // the caller's float declare reads the rebuilt weight + } + + r.w = L.typed(GGML_TYPE_I8, "%s.weight", site); + r.wscale = L.f32("%s.wscale", site); + r.ascale = L.opt_typed(GGML_TYPE_F32, "%s.ascale", site); + r.bias = has_bias ? L.f32("%s.bias", site) : nullptr; + if (!r.w || !r.wscale || (has_bias && !r.bias)) { + r.w = nullptr; + return r; + } + if (!fill_specs(L, mod, site_key, site, gamma, eps, r)) + r.w = nullptr; + return r; +} + +void fq_set_heads(FqLinear & s, int head_dim, int heads, uint32_t vmask) { + static const bool off = [] { const char * e = std::getenv("VLA_FQ_NO_HEADS"); return e && *e && *e != '0'; }(); + if (!s.w || off || s.heads_off) return; + const int64_t parts = s.gemm.N / ((int64_t) head_dim * heads); + if (parts < 1 || parts > 32 || s.gemm.N % ((int64_t) head_dim * heads) != 0) return; + s.gemm.head_dim = head_dim; + s.gemm.heads = heads; + s.gemm.vmask = vmask; +} + +void fq_link_prefetch(const std::vector & order) { + FqLinear * prev = nullptr; + for (FqLinear * s : order) { + if (!s || !s->w) continue; + if (prev) prev->gemm.next_w = s->w; + prev = s; + } +} + +FqLinear fq_declare_fused(WeightLoader & L, const FqModuleSpec & mod, const char * site_key, + bool has_bias, const std::string & out_base, const std::vector & sites) { + FqLinear r; + if (sites.empty()) + return r; + + std::vector ws, ss, bs; + for (const std::string & s : sites) { + const std::string wname = s + ".weight"; + const ggml_tensor * meta = L.reader().meta(wname.c_str()); + if (!meta || meta->type != GGML_TYPE_I8) { + if (&s != &sites[0]) { + std::fprintf(stderr, "vla: %s: FoldQuant sites fused as %s must all be INT8 or all float\n", + s.c_str(), out_base.c_str()); + L.fail("FoldQuant fused group"); + } + return r; + } + ws.push_back(wname); + ss.push_back(s + ".wscale"); + bs.push_back(s + ".bias"); + } + if (mod.dequant) { + // Each site's ascale goes into its own rows, so they need not agree here. + for (const std::string & s : sites) + if (!register_dequant(L, mod, site_key, s)) + break; + return r; // the caller's fuse_gemm reads the rebuilt weights + } + + r.w = L.fuse_typed(GGML_TYPE_I8, (out_base + ".w").c_str(), ws); + r.wscale = L.fuse_f32((out_base + ".wscale").c_str(), ss); + r.bias = has_bias ? L.fuse_f32((out_base + ".b").c_str(), bs) : nullptr; + if (!r.w || !r.wscale || (has_bias && !r.bias)) { + r.w = nullptr; + return r; + } + + // The fused group shares one input transform, so the SQ vectors must agree. + const std::string a0 = sites[0] + ".ascale"; + const bool has_a0 = L.reader().meta(a0.c_str()) != nullptr; + for (size_t i = 1; i < sites.size(); ++i) { + const std::string ai = sites[i] + ".ascale"; + const bool has_ai = L.reader().meta(ai.c_str()) != nullptr; + if (has_ai != has_a0) { + std::fprintf(stderr, "vla: %s: ascale present on some fused sites but not others\n", out_base.c_str()); + L.fail("FoldQuant fused ascale"); + r.w = nullptr; + return r; + } + if (has_a0) { + const std::vector v0 = L.reader().read_f32(a0.c_str()); + const std::vector vi = L.reader().read_f32(ai.c_str()); + if (v0.empty() || v0 != vi) { + std::fprintf(stderr, "vla: %s and %s differ; fused sites must share one SQ vector\n", + a0.c_str(), ai.c_str()); + L.fail("FoldQuant fused ascale"); + r.w = nullptr; + return r; + } + } + } + r.ascale = has_a0 ? L.typed(GGML_TYPE_F32, "%s", a0.c_str()) : nullptr; + + if (!fill_specs(L, mod, site_key, out_base.c_str(), nullptr, 0.0f, r)) + r.w = nullptr; + return r; +} + +} // namespace vla diff --git a/src/foldquant.h b/src/foldquant.h new file mode 100644 index 0000000..28a87ba --- /dev/null +++ b/src/foldquant.h @@ -0,0 +1,179 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant: real-quantized linears (docs/QUANTIZATION.md). +// +// A FoldQuant site ships `.weight` as GGML_TYPE_I8 codes (per-output-row +// symmetric, in a block-Hadamard-rotated frame), `.wscale` F32[N], and +// optionally `.ascale` F32[K] (the static SmoothQuant vector). At runtime +// the activation is [RMS-normed,] [divided,] rotated by the same butterfly, +// quantized per token to INT8, multiplied on integer units and dequantized. +// +// Every backend that runs a site natively runs the same two GGML_OP_CUSTOM nodes +// (fq_act, fq_gemm; see layers/fq_linear.h): the CPU backend executes the +// reference in foldquant_ref.h, CUDA and SYCL claim them through their ggml +// extension hooks, and OpenVINO's CPU and GPU plugins translate them into +// OpenVINO ops (src/openvino/foldquant_ov.cpp). +// +// The other backends (Metal, Vulkan, Hexagon, OpenCL, the OpenVINO NPU) have no +// implementation of those nodes, so there a site is read back as a float GEMM +// weight instead (dequant mode): W_deq . B . diag(1/ascale) with the ascale +// divide on the side of B the activation applies it, rebuilt at upload in the +// resident GEMM type. The arch then takes its stock float path, with the folded +// norm gains the file already carries. The weights keep FoldQuant's rounding; +// the activations stay float (no per-token quantization, no INT4 clip). + +#pragma once + +#include "ggml.h" + +#include +#include +#include +#include + +namespace vla { + +struct gguf_reader; +struct Backend; +class WeightLoader; + +constexpr uint32_t FQ_ACT_MAGIC = 0x31414651u; // "FQA1" +constexpr uint32_t FQ_GEMM_MAGIC = 0x31474651u; // "FQG1" + +// Bytes after the codes in every row of the activation blob: the per-token +// float scale sits at byte K_pack, the rest is padding that keeps rows 16-byte +// aligned (K is a multiple of 64 wherever a site is accepted). +constexpr int64_t FQ_ACT_TAIL = 16; + +// Per-module scheme parameters parsed from `.quant.*`. +struct FqModuleSpec { + int wbits = 8; // 8 | 4 + int abits = 8; // 8 | 4 + int rot_block = 64; // nominal; narrowed per site by fq_rot_block_for() + bool fold_before = false; // ascale divides before (true) or after the butterfly + float clip = 1.0f; // act_clip_ratio (INT4 activations only) + std::map site_bits; // per-site override, e.g. {"o":8,"down":8} + std::string scheme; + bool dequant = false; // read every site back as a float weight (see above) + bool no_heads = false; // keep plain [N][T] GEMM outputs (no fq_set_heads layout) + + int wbits_for(const char * site_key) const; + int abits_for(const char * site_key) const; +}; + +struct FoldQuantSpec { + bool present = false; + FqModuleSpec llm, action; + std::string method, applied_at, provenance; +}; + +// Static per-node parameters. Graph nodes reference them by pointer (custom-op +// userdata), so they must outlive every graph: they live inside the module +// weight structs, which are sized once at declare time and never resized. +struct FqActSpec { + uint32_t magic = FQ_ACT_MAGIC; + int64_t K = 0; + int abits = 8; + int rot_block = 1; // 1 = no rotation + bool fold_before = false; + bool has_gamma = false; // fused RMSNorm with the folded gamma + bool has_ascale = false; // static SmoothQuant vector shipped with the site + float clip = 1.0f; + float eps = 0.0f; // RMSNorm epsilon when has_gamma +}; + +struct FqGemmSpec { + uint32_t magic = FQ_GEMM_MAGIC; + int64_t K = 0; + int64_t N = 0; + int wbits = 8; + // The weight tensor the next FoldQuant GEMM in execution order reads; the + // CUDA kernel prefetches its head into L2 once its own loads are issued + // (fq_link_prefetch; opt-in via VLA_FQ_PREFETCH_MB, measured slower on Orin). + const ggml_tensor * next_w = nullptr; + // Head layout of the output (fq_set_heads): the N columns are nparts + // consecutive [heads][head_dim] projections and the epilogue writes each + // part directly in the layout attention consumes, [head_dim][T][heads] + // (Q/K) or [T][head_dim][heads] (parts whose bit is set in vmask), so the + // permute copies after the GEMM disappear. 0 = plain [N][T]. + int head_dim = 0; + int heads = 0; + uint32_t vmask = 0; +}; + +struct FqLinear { + ggml_tensor * w = nullptr; // I8 [K_pack, N]; null => not a FoldQuant site + ggml_tensor * wscale = nullptr; // F32 [N] + ggml_tensor * ascale = nullptr; // F32 [K] or null + ggml_tensor * gamma = nullptr; // F32 [K] folded RMSNorm gamma or null + ggml_tensor * bias = nullptr; // F32 [N] or null + FqActSpec act; + FqGemmSpec gemm; + bool heads_off = false; // fq_set_heads is a no-op (the backend reads no raw layout) + + explicit operator bool() const { return w != nullptr; } +}; + +// Activation blob layout shared by the CPU reference, the CUDA kernels and the +// tests: I8 tensor, ne = [row_bytes, T]; row t = codes[0, K_pack) then float scale. +inline int64_t fq_act_kpack(int64_t K, int abits) { return abits == 4 ? K / 2 : K; } +inline int64_t fq_act_row_bytes(int64_t K, int abits) { return fq_act_kpack(K, abits) + FQ_ACT_TAIL; } +inline int64_t fq_w_kpack(int64_t K, int wbits) { return wbits == 4 ? K / 2 : K; } + +// Largest power of two <= nominal that divides K (FoldQuant's rotation_block_for); +// 1 means no rotation. +int fq_rot_block_for(int64_t K, int nominal); + +// `.quant.method` == "foldquant". Never asserts on a wrong KV type. +bool foldquant_present(const gguf_reader & g, const char * prefix); +FoldQuantSpec foldquant_parse (const gguf_reader & g, const char * prefix); + +// Load-time policy: CUDA and SYCL register their integer kernels, CPU runs the +// exact reference, OpenVINO (CPU and GPU devices) translates the nodes into +// OpenVINO ops (src/openvino/foldquant_ov.cpp); any other backend, the OpenVINO +// NPU, and VLA_FQ_DEQUANT=1 switch fq to dequant mode. Returns false only on an unusable spec. +bool foldquant_check_backend(const char * tag, const Backend & b, FoldQuantSpec & fq, bool weight_dtype_set); + +// A site's float weight in dequant mode, row-major [N][K]: per row, the codes +// times wscale, then the ascale divide and the block rotation in the order the +// activation applies them reversed onto the weight. codes: [N][K] int8 for W8, +// [N][K/2] nibbles (low nibble = even column) for W4. ascale may be null. +void fq_dequant_rows(const int8_t * codes, const float * wscale, const float * ascale, + int64_t K, int64_t N, int wbits, int rot_block, bool fold_before, float * out); + +// Declares `.weight` (I8), `.wscale`, optional `.ascale` and optional +// `.bias`. Returns an empty FqLinear (w == nullptr) when `.weight` is not +// I8, so the caller falls back to its stock declare. In dequant mode it declares +// nothing: it registers `.weight` with the loader as a float weight and +// returns an empty FqLinear, so the caller's stock declare reads it back. `gamma` is the already +// declared folded norm weight fused into the activation node, or null. +FqLinear fq_declare_linear(WeightLoader & L, const FqModuleSpec & mod, const char * site_key, + bool has_bias, ggml_tensor * gamma, float eps, + const char * site_fmt, ...) __attribute__((format(printf, 7, 8))); + +// Several sites sharing one input transform (DiT q/k/v, k/v) fused into one +// weight: codes and wscale concatenate along N; ascale must agree and is kept once. +// Marks a site's output as head-laid-out (see FqGemmSpec). No-op when the +// site is float or VLA_FQ_NO_HEADS=1 (A/B switch: identical numbers either way). +void fq_set_heads(FqLinear & s, int head_dim, int heads, uint32_t vmask); + +// Chains sites in execution order so each GEMM knows the next site's weights +// (null entries and absent sites are skipped). +void fq_link_prefetch(const std::vector & order); + +FqLinear fq_declare_fused(WeightLoader & L, const FqModuleSpec & mod, const char * site_key, + bool has_bias, const std::string & out_base, const std::vector & sites); + +} // namespace vla diff --git a/src/foldquant_ref.cpp b/src/foldquant_ref.cpp new file mode 100644 index 0000000..7f0698b --- /dev/null +++ b/src/foldquant_ref.cpp @@ -0,0 +1,303 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant reference bodies. This translation unit is compiled with +// -ffp-contract=off (CMakeLists.txt); see foldquant_ref.h for why that matters. + +#include "foldquant_ref.h" +#include "env_flag.h" + +#include +#include +#include +#include +#include +#include + +namespace vla { + +static bool fq_stats_enabled() { + static const bool on = env_flag("VLA_FQ_STATS"); + return on; +} + +namespace fqref { + +// The kernel runs one warp per row: lane l owns the 64-element chunks +// c = l, l+32, l+64, ... of the row (K is a multiple of 64), accumulates its +// partial sum of squares over them in element order, and the 32 partials are +// combined by an xor-shuffle butterfly. NT is the number of partials. + +// Kernel-order reduction of the 32 lane partials. Clobbers p. +float block_sum(float * p) { + float t[32]; + for (int off = 16; off > 0; off >>= 1) { + for (int l = 0; l < 32; ++l) t[l] = p[l] + p[l ^ off]; + std::memcpy(p, t, sizeof(t)); + } + return p[0]; +} + +// In-place natural-order Sylvester-Hadamard butterfly on every bs-block of a +// row, normalised by inv_sqrt_bs = 1/sqrt(bs) afterwards (same stage order and +// pairing as FoldQuant's block FWHT kernel). +void fwht_row(float * y, int64_t K, int bs, float inv_sqrt_bs) { + for (int64_t b = 0; b < K; b += bs) { + float * v = y + b; + for (int h = 1; h < bs; h <<= 1) + for (int i = 0; i < bs; i += 2*h) + for (int j = i; j < i + h; ++j) { + const float a = v[j], c = v[j + h]; + v[j] = a + c; + v[j + h] = a - c; + } + for (int j = 0; j < bs; ++j) v[j] = v[j] * inv_sqrt_bs; + } +} + +float inv_sqrt_block(int bs) { + return (float) (1.0 / std::sqrt((double) bs)); +} + +float qmax_for(int bits) { + return bits == 4 ? 7.0f : 127.0f; +} + +// One activation row: x[K] -> codes + scale. tmp and partial are scratch +// (K floats and NT floats). +void act_row(const float * x, const float * ascale, const float * gamma, const FqActSpec & s, + float * tmp, float * partial, int8_t * codes_out, float * scale_out) { + const int64_t K = s.K; + + if (s.has_gamma && gamma) { + for (int t = 0; t < NT; ++t) partial[t] = 0.0f; + for (int64_t k = 0; k < K; ++k) { + const float v = x[k]; + const int lane = (int) ((k / CHUNK) % NT); + partial[lane] = partial[lane] + v * v; + } + const float sumsq = block_sum(partial); + const float rstd = 1.0f / std::sqrt(sumsq / (float) K + s.eps); + for (int64_t k = 0; k < K; ++k) tmp[k] = (x[k] * rstd) * gamma[k]; + } else { + for (int64_t k = 0; k < K; ++k) tmp[k] = x[k]; + } + + if (ascale && s.fold_before) + for (int64_t k = 0; k < K; ++k) tmp[k] = tmp[k] / ascale[k]; + + if (s.rot_block > 1) + fwht_row(tmp, K, s.rot_block, inv_sqrt_block(s.rot_block)); + + if (ascale && !s.fold_before) + for (int64_t k = 0; k < K; ++k) tmp[k] = tmp[k] / ascale[k]; + + float amax = 0.0f; + for (int64_t k = 0; k < K; ++k) amax = std::max(amax, std::fabs(tmp[k])); + + const float qmax = qmax_for(s.abits); + float scale = (s.clip * amax) / qmax; + if (scale < 1e-12f) scale = 1e-12f; + *scale_out = scale; + // Multiply by the reciprocal rather than divide: what FoldQuant's TensorRT + // kernels do, so the deployed engines and this runtime round alike. + const float inv = 1.0f / scale; + + if (s.abits == 4) { + for (int64_t k = 0; k < K; k += 2) { + float q0 = std::nearbyint(tmp[k] * inv); + float q1 = std::nearbyint(tmp[k + 1] * inv); + q0 = std::min(qmax, std::max(-qmax, q0)); + q1 = std::min(qmax, std::max(-qmax, q1)); + const int i0 = (int) q0, i1 = (int) q1; + codes_out[k / 2] = (int8_t) ((i0 & 0xF) | ((i1 & 0xF) << 4)); + } + } else { + for (int64_t k = 0; k < K; ++k) { + float q = std::nearbyint(tmp[k] * inv); + q = std::min(qmax, std::max(-qmax, q)); + codes_out[k] = (int8_t) (int) q; + } + } +} + +// Unpack a nibble row (low nibble = even column, two's complement) to int8. +void unpack_nibbles(const int8_t * packed, int64_t K, int8_t * out) { + for (int64_t k = 0; k < K; k += 2) { + const uint8_t b = (uint8_t) packed[k / 2]; + out[k] = (int8_t) ((int8_t) (b << 4) >> 4); + out[k + 1] = (int8_t) ((int8_t) b >> 4); + } +} + +// y[n] for one (token, output row): exact int32 accumulation, then the +// kernel's epilogue order. +float gemm_dot(const int8_t * w_row, const int8_t * x_row, int64_t K, float xs, float ws, float bias) { + int32_t acc = 0; + for (int64_t k = 0; k < K; ++k) acc += (int32_t) w_row[k] * (int32_t) x_row[k]; + float v = ((float) acc * xs) * ws; + return v + bias; +} + +} // namespace fqref + +// GGML_OP_CUSTOM entry points. Sources are packed without holes: +// fq_act : src[0]=x F32 [K, T...] contiguous, then gamma F32[K] if has_gamma, +// then ascale F32[K] if has_ascale; dst I8 [row_bytes, T] +// fq_gemm: src[0]=w I8 [K_pack, N], src[1]=xq blob, src[2]=wscale F32[N], src[3]=bias F32[N] or null, +// src[4]=residual F32 [N, T] or absent (the residual add fused into the epilogue); dst F32 [N, T] + +void fq_act_cpu(ggml_tensor * dst, int ith, int nth, void * userdata) { + const FqActSpec & s = *(const FqActSpec *) userdata; + if (s.magic != FQ_ACT_MAGIC) return; + + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * g = nullptr, * as = nullptr; + fq_act_srcs(dst, s, &g, &as); + + const int64_t K = s.K; + const int64_t T = dst->ne[1]; + const int64_t rb = fq_act_row_bytes(K, s.abits); + const int64_t kp = fq_act_kpack(K, s.abits); + const int64_t per = (T + nth - 1) / nth; + + // VLA_FQ_STATS=1: one line per node with checksums of every input, so two + // backends running this same reference can be diffed node by node. + if (ith == 0 && fq_stats_enabled()) { + double sum = 0.0, sumsq = 0.0, mx = 0.0; + for (int64_t t = 0; t < T; ++t) { + const float * xr = (const float *) ((const char *) x->data + t * x->nb[1]); + for (int64_t k = 0; k < K; ++k) { const double v = xr[k]; sum += v; sumsq += v*v; mx = std::fmax(mx, std::fabs(v)); } + } + double gs = 0.0, as_ = 0.0; + if (g) for (int64_t k = 0; k < K; ++k) gs += ((const float *) g->data)[k]; + if (as) for (int64_t k = 0; k < K; ++k) as_ += ((const float *) as->data)[k]; + std::fprintf(stderr, "FQSTAT act %-44s T=%lld x sum=%.9g sumsq=%.9g max=%.9g gamma=%.9g ascale=%.9g\n", + ggml_get_name(dst), (long long) T, sum, sumsq, mx, gs, as_); + // VLA_FQ_DUMP=

: also write this node's input rows as raw F32 so two + // backends can be compared element by element (first occurrence only). + if (const char * dir = std::getenv("VLA_FQ_DUMP")) { + static int n_dumped = 0; + if (n_dumped < 64) { + char path[1024]; + std::snprintf(path, sizeof path, "%s/%s.%d.f32", dir, ggml_get_name(dst), n_dumped++); + if (FILE * f = std::fopen(path, "wb")) { + for (int64_t t = 0; t < T; ++t) + std::fwrite((const char *) x->data + t * x->nb[1], sizeof(float), (size_t) K, f); + std::fclose(f); + } + } + } + } + const int64_t t0 = (int64_t) ith * per; + const int64_t t1 = std::min(T, t0 + per); + + const float * xp = (const float *) x->data; + const float * ap = as ? (const float *) as->data : nullptr; + const float * gp = g ? (const float *) g->data : nullptr; + + std::vector tmp((size_t) K); + float partial[fqref::NT]; + for (int64_t t = t0; t < t1; ++t) { + uint8_t * row = (uint8_t *) dst->data + (size_t) t * rb; + fqref::act_row((const float *) ((const char *) xp + (size_t) t * x->nb[1]), ap, gp, s, tmp.data(), partial, + (int8_t *) row, (float *) (row + kp)); + } +} + +// Where output element (row t, column n) lands: plain row-major, or the head +// layout described in FqGemmSpec (same rule as the CUDA epilogues). +static inline size_t fq_out_index(const FqGemmSpec & s, int64_t T, int64_t t, int64_t n) { + if (!s.heads) return (size_t) t * s.N + n; + const int64_t dim = (int64_t) s.head_dim * s.heads; + const int64_t p = n / dim, r = n - p * dim, h = r / s.head_dim, d = r - h * s.head_dim; + const int64_t base = p * dim * T; + return (size_t) (((s.vmask >> p) & 1) ? base + (h * s.head_dim + d) * T + t + : base + (h * T + t) * s.head_dim + d); +} + +void fq_gemm_cpu(ggml_tensor * dst, int ith, int nth, void * userdata) { + const FqGemmSpec & s = *(const FqGemmSpec *) userdata; + if (s.magic != FQ_GEMM_MAGIC) return; + + const ggml_tensor * w = dst->src[0]; + const ggml_tensor * xq = dst->src[1]; + const ggml_tensor * ws = dst->src[2]; + const ggml_tensor * b = dst->src[3]; + const ggml_tensor * r = dst->src[4]; // optional residual [N, T] + + const int64_t K = s.K; + const int64_t N = s.N; + const int64_t T = dst->ne[1]; + const int abits = (xq->ne[0] == K + FQ_ACT_TAIL) ? 8 : 4; + const int64_t rb = fq_act_row_bytes(K, abits); + const int64_t kp_a = fq_act_kpack(K, abits); + const int64_t kp_w = fq_w_kpack(K, s.wbits); + + if (ith == 0 && fq_stats_enabled()) { + long long csum = 0; double ssum = 0.0; + for (int64_t t = 0; t < T; ++t) { + const int8_t * row = (const int8_t *) xq->data + t * rb; + for (int64_t k = 0; k < kp_a; ++k) csum += row[k]; + float sc; std::memcpy(&sc, row + kp_a, sizeof sc); ssum += sc; + } + long long wsum = 0; + for (int64_t i = 0; i < (int64_t) N * kp_w; ++i) wsum += ((const int8_t *) w->data)[i]; + double wss = 0.0, bs = 0.0; + for (int64_t n = 0; n < N; ++n) wss += ((const float *) ws->data)[n]; + if (b) for (int64_t n = 0; n < N; ++n) bs += ((const float *) b->data)[n]; + std::fprintf(stderr, "FQSTAT gemm %-44s T=%lld codes=%lld scales=%.9g w=%lld wscale=%.9g bias=%.9g\n", + ggml_get_name(dst), (long long) T, csum, ssum, wsum, wss, bs); + } + + const int8_t * wp = (const int8_t *) w->data; + const uint8_t * xp = (const uint8_t *) xq->data; + const float * wsp = (const float *) ws->data; + const float * bp = b ? (const float *) b->data : nullptr; + const float * rp = r ? (const float *) r->data : nullptr; + float * y = (float *) dst->data; + + // Activations unpacked once per thread when they are nibbles. + std::vector xa; + if (abits == 4) { + xa.resize((size_t) T * K); + for (int64_t t = 0; t < T; ++t) + fqref::unpack_nibbles((const int8_t *) (xp + (size_t) t * rb), K, xa.data() + (size_t) t * K); + } + std::vector wrow_buf(s.wbits == 4 ? (size_t) K : 0); + + const int64_t per = (N + nth - 1) / nth; + const int64_t n0 = (int64_t) ith * per; + const int64_t n1 = std::min(N, n0 + per); + for (int64_t n = n0; n < n1; ++n) { + const int8_t * wrow = wp + (size_t) n * kp_w; + if (s.wbits == 4) { + fqref::unpack_nibbles(wrow, K, wrow_buf.data()); + wrow = wrow_buf.data(); + } + const float wsn = wsp[n]; + const float bn = bp ? bp[n] : 0.0f; + for (int64_t t = 0; t < T; ++t) { + const uint8_t * row = xp + (size_t) t * rb; + const int8_t * xrow = (abits == 4) ? xa.data() + (size_t) t * K : (const int8_t *) row; + float xs; + std::memcpy(&xs, row + kp_a, sizeof(float)); + float v = fqref::gemm_dot(wrow, xrow, K, xs, wsn, bn); + if (rp) v = v + rp[(size_t) t * N + n]; + y[fq_out_index(s, T, t, n)] = v; + } + } +} + +} // namespace vla diff --git a/src/foldquant_ref.h b/src/foldquant_ref.h new file mode 100644 index 0000000..91337e6 --- /dev/null +++ b/src/foldquant_ref.h @@ -0,0 +1,87 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// The FoldQuant reference: what the CPU backend runs, and what the CUDA kernels +// must reproduce bit for bit. Every float operation is written in the order the +// kernels perform it (per-thread strided partial sums, a fixed reduction tree, +// IEEE sqrt/division, round-half-even), and the translation units that include +// this header are compiled with FP contraction off, so a+b*c never becomes an +// FMA on one side only. +// +// The bodies live in foldquant_ref.cpp, the one translation unit compiled with +// -ffp-contract=off: as inline functions in a header they were compiled once per +// including file, most of them with the compiler's default FMA contraction, and +// the linker kept an arbitrary copy - a reference that rounded differently from +// the kernels on every bias add and every sum of squares. + +#pragma once + +#include "foldquant.h" + +#include "ggml.h" + +#include +#include +#include +#include +#include + +namespace vla { +namespace fqref { + +// "Threads" per row in the kernel: lane l owns the 64-element chunks +// c = l, l+32, ... of the row and accumulates its partial sum of squares over +// them in element order; the 32 partials are combined by an xor butterfly. +constexpr int NT = 32; +constexpr int CHUNK = 64; + +// Kernel-order reduction of the 32 lane partials. Clobbers p. +float block_sum(float * p); + +// In-place natural-order Sylvester-Hadamard butterfly on every bs-block of a +// row, normalised by inv_sqrt_bs = 1/sqrt(bs) afterwards. +void fwht_row(float * y, int64_t K, int bs, float inv_sqrt_bs); +float inv_sqrt_block(int bs); +float qmax_for(int bits); + +// One activation row: x[K] -> codes + scale. tmp and partial are scratch +// (K floats and NT floats). +void act_row(const float * x, const float * ascale, const float * gamma, const FqActSpec & s, + float * tmp, float * partial, int8_t * codes_out, float * scale_out); + +// Unpack a nibble row (low nibble = even column, two's complement) to int8. +void unpack_nibbles(const int8_t * packed, int64_t K, int8_t * out); + +// y for one (token, output row): exact int32 accumulation, then the kernel's +// epilogue order ((float)acc * xs) * ws + bias. +float gemm_dot(const int8_t * w_row, const int8_t * x_row, int64_t K, float xs, float ws, float bias); + +} // namespace fqref + +// GGML_OP_CUSTOM entry points. Sources are packed without holes: +// fq_act : src[0]=x F32 [K, T...] contiguous, then gamma F32[K] if has_gamma, +// then ascale F32[K] if has_ascale; dst I8 [row_bytes, T] +// fq_gemm: src[0]=w I8 [K_pack, N], src[1]=xq blob, src[2]=wscale F32[N], src[3]=bias F32[N] or null +// dst F32 [N, T] +inline void fq_act_srcs(const ggml_tensor * dst, const FqActSpec & s, + const ggml_tensor ** gamma, const ggml_tensor ** ascale) { + int i = 1; + *gamma = s.has_gamma ? dst->src[i++] : nullptr; + *ascale = s.has_ascale ? dst->src[i++] : nullptr; +} + +void fq_act_cpu (ggml_tensor * dst, int ith, int nth, void * userdata); +void fq_gemm_cpu(ggml_tensor * dst, int ith, int nth, void * userdata); + +} // namespace vla diff --git a/src/kernels/foldquant/fq_gemm_i8.cu b/src/kernels/foldquant/fq_gemm_i8.cu new file mode 100644 index 0000000..b3fa656 --- /dev/null +++ b/src/kernels/foldquant/fq_gemm_i8.cu @@ -0,0 +1,392 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// INT8 x INT8 -> INT32 GEMM on the integer tensor cores (wmma s8, 16x16x16), +// K streamed through a multi-stage cp.async pipeline so that while one chunk +// of the weight stripe is being multiplied the next ones are in flight: at +// the small M a VLA runs (a few hundred tokens, 41 action tokens) the GEMM is +// bound by streaming W, and the pipeline is what keeps the DRAM busy. +// +// Four tilings, picked per shape by launch_gemm (see there) from a measured +// sweep on a 16-SM Orin: narrow-N sites at M <= 64 get a 32-wide N tile (CTA +// count), the LLM prefill a 192-row M tile that holds ~160 tokens in one tile, +// and the widest / longest sites a 128-row BK=128 tile. VLA_FQ_TILE= +// forces one of the twelve candidates for re-tuning on another GPU. +// +// A = the activation blob (row m at m*row_bytes, codes then the float scale), +// B = W[N][K] with K contiguous, which is exactly a col_major K x N operand. +// Fused dequant epilogue y = ((float)acc * act_scale[m]) * wscale[n] (+ bias[n]) +// in the order foldquant_ref.h uses; the accumulation is integer so any tiling +// gives the same result. +// +// M tails are zero-filled; N and K are multiples of 64 (checked at load). + +#include "fq_kernels.h" + +#include + +#include +#include + +namespace vla { +namespace fq { + +namespace { + +using namespace nvcuda; + +__device__ __forceinline__ void cp_async16(void * smem, const void * gmem, bool pred) { + const unsigned s = (unsigned) __cvta_generic_to_shared(smem); + const int sz = pred ? 16 : 0; // src-size 0 zero-fills the 16 bytes + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" :: "r"(s), "l"(gmem), "r"(sz)); +} +__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); } +template +__device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); } + +template +struct Tile { + static constexpr int THREADS = WM * WN * 32; + static constexpr int LDS = BK + 16; // smem row stride in bytes: 16-byte aligned rows, fewer bank conflicts + static constexpr int LDC = BN + 4; + static constexpr int STAGE_BYTES = (BM + BN) * LDS; + static constexpr int C_BYTES = BM * LDC * 4; + // The int32 C staging reuses the drained pipeline buffers. + static constexpr int SMEM_BYTES = STAGES * STAGE_BYTES > C_BYTES ? STAGES * STAGE_BYTES : C_BYTES; + static constexpr int CHUNKS = (BM + BN) * (BK / 16); // 16-byte copies per stage + static constexpr int PER_THR = CHUNKS / THREADS; + static constexpr int FM = BM / WM / 16, FN = BN / WN / 16; // 16x16 fragments per warp + static_assert(CHUNKS % THREADS == 0, "tile copies must split evenly over the threads"); + static_assert(BM % (WM * 16) == 0 && BN % (WN * 16) == 0, "warp tile must be a multiple of 16"); +}; + +// SPLITK > 1: gridDim.z CTAs each accumulate one K slice of the tile into +// registers, add their int32 partials into a zeroed workspace with integer +// atomics (exact and associative, so the sum does not depend on arrival +// order) and the last CTA to arrive - counted per tile - applies the float +// epilogue and re-zeroes the workspace for the next launch. Streams each +// weight byte once, like SPLITK == 1, but with gridDim.z times the CTAs in +// flight, which is what a DRAM-bound M <= 64 site needs. +template +__global__ void __launch_bounds__(WM * WN * 32) gemm_i8_kernel(const GemmArgs g) { + using T = Tile; + // Dynamic: the wide tiles need more than the 48 KB static limit (Orin + // allows 164 KB per SM, opted into once per instantiation in launch_tile). + extern __shared__ __align__(128) int8_t smem[]; + int * Cs = (int *) smem; + + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + // M tiles are the fastest grid axis: the CTAs that share a weight stripe + // run back to back and the stripe is served from L2 after the first. + const int64_t m0 = (int64_t) blockIdx.x * BM; + const int64_t n0 = (int64_t) blockIdx.y * BN; + const int KT_all = (int) (g.K / BK); + const int KT = SPLITK > 1 ? KT_all / SPLITK : KT_all; // this CTA's k-steps + const int kt_base = SPLITK > 1 ? (int) blockIdx.z * KT : 0; + + // Per-thread copy list for one stage: global source (at k0 = 0), smem + // offset and the zero-fill predicate for A rows past M. + const int8_t * src[T::PER_THR]; + int dst[T::PER_THR]; + bool pred[T::PER_THR]; + #pragma unroll + for (int i = 0; i < T::PER_THR; ++i) { + const int c = tid + i * T::THREADS; + const int r = c / (BK / 16), c16 = c % (BK / 16); + if (r < BM) { + const int64_t m = m0 + r; + pred[i] = m < g.M; + src[i] = g.blob + (pred[i] ? m : 0) * g.row_bytes + c16 * 16; + dst[i] = r * T::LDS + c16 * 16; + } else { + const int rb = r - BM; + pred[i] = true; + src[i] = g.w + (n0 + rb) * g.K + c16 * 16; + dst[i] = BM * T::LDS + rb * T::LDS + c16 * 16; + } + } + auto load_tile = [&](int stage, int kt) { + int8_t * base = smem + stage * T::STAGE_BYTES; + const int64_t k0 = (int64_t) (kt_base + kt) * BK; + #pragma unroll + for (int i = 0; i < T::PER_THR; ++i) cp_async16(base + dst[i], src[i] + k0, pred[i]); + }; + + wmma::fragment acc[T::FM][T::FN]; + #pragma unroll + for (int i = 0; i < T::FM; ++i) + #pragma unroll + for (int j = 0; j < T::FN; ++j) + wmma::fill_fragment(acc[i][j], 0); + + #pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < KT) load_tile(s, s); + cp_async_commit(); + } + + for (int kt = 0; kt < KT; ++kt) { + cp_async_wait(); + __syncthreads(); // stage kt%STAGES landed for every thread; stage (kt-1)%STAGES is free + const int stage = kt % STAGES; + const int nk = kt + STAGES - 1; + if (nk < KT) load_tile(nk % STAGES, nk); + cp_async_commit(); + + const int8_t * as = smem + stage * T::STAGE_BYTES; + const int8_t * bs = as + BM * T::LDS; + #pragma unroll + for (int kk = 0; kk < BK; kk += 16) { + wmma::fragment af[T::FM]; + wmma::fragment bf[T::FN]; + #pragma unroll + for (int i = 0; i < T::FM; ++i) + wmma::load_matrix_sync(af[i], as + (wm * T::FM * 16 + i * 16) * T::LDS + kk, T::LDS); + #pragma unroll + for (int j = 0; j < T::FN; ++j) + wmma::load_matrix_sync(bf[j], bs + (wn * T::FN * 16 + j * 16) * T::LDS + kk, T::LDS); + #pragma unroll + for (int i = 0; i < T::FM; ++i) + #pragma unroll + for (int j = 0; j < T::FN; ++j) + wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]); + } + } + cp_async_wait<0>(); + __syncthreads(); + + #pragma unroll + for (int i = 0; i < T::FM; ++i) + #pragma unroll + for (int j = 0; j < T::FN; ++j) + wmma::store_matrix_sync(Cs + (wm * T::FM * 16 + i * 16) * T::LDC + wn * T::FN * 16 + j * 16, acc[i][j], + T::LDC, wmma::mem_row_major); + __syncthreads(); + + if (SPLITK > 1) { + // Partial tile -> workspace (int32 [M][N]); the last CTA of this tile finishes. + __shared__ int s_last; + for (int idx = tid; idx < BM * BN; idx += T::THREADS) { + const int r = idx / BN, c = idx - r * BN; + const int64_t m = m0 + r, n = n0 + c; + if (m >= g.M) continue; + atomicAdd(g.ws + m * g.N + n, Cs[r * T::LDC + c]); + } + __threadfence(); + __syncthreads(); + if (tid == 0) { + const int tile = (int) (blockIdx.y * gridDim.x + blockIdx.x); + s_last = (atomicAdd(g.counters + tile, 1) == SPLITK - 1); + } + __syncthreads(); + if (!s_last) return; + __threadfence(); + for (int idx = tid; idx < BM * BN; idx += T::THREADS) { + const int r = idx / BN, c = idx - r * BN; + const int64_t m = m0 + r, n = n0 + c; + if (m >= g.M) continue; + int * wp = g.ws + m * g.N + n; + const int acc = __ldcg(wp); + *wp = 0; // ready for the next launch + const float xs = *(const float *) (g.blob + m * g.row_bytes + g.K); + float v = ((float) acc * xs) * g.wscale[n]; + if (HAS_BIAS) v = v + g.bias[n]; + if (g.res) v = v + g.res[m * g.N + n]; + g.y[fq_out_index(g, m, n)] = v; + } + if (tid == 0) g.counters[(int) (blockIdx.y * gridDim.x + blockIdx.x)] = 0; + return; + } + + for (int idx = tid; idx < BM * BN; idx += T::THREADS) { + const int r = idx / BN, c = idx - r * BN; + const int64_t m = m0 + r, n = n0 + c; + if (m >= g.M) continue; + const float xs = *(const float *) (g.blob + m * g.row_bytes + g.K); + float v = ((float) Cs[r * T::LDC + c] * xs) * g.wscale[n]; + if (HAS_BIAS) v = v + g.bias[n]; + if (g.res) v = v + g.res[m * g.N + n]; + g.y[fq_out_index(g, m, n)] = v; + } +} + +// Split-K workspace: one int32 [M][N] accumulator plus one counter per tile, +// zero between launches (the finishing CTA re-zeroes what it used). GEMMs on a +// stream run in order, so one buffer serves them all; it grows on demand, +// never inside CUDA-graph capture (the caller then gets cudaErrorNotSupported +// and falls back to the unsplit tile). +struct SplitWs { + int * ws = nullptr; + int * counters = nullptr; + size_t ws_elems = 0, n_counters = 0; +}; +static SplitWs & split_ws() { static SplitWs w; return w; } + +static cudaError_t ensure_split_ws(size_t elems, size_t counters, cudaStream_t stream) { + SplitWs & w = split_ws(); + if (w.ws_elems >= elems && w.n_counters >= counters) return cudaSuccess; + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + if (cudaStreamIsCapturing(stream, &cap) == cudaSuccess && cap != cudaStreamCaptureStatusNone) + return cudaErrorNotSupported; + if (cudaStreamSynchronize(stream) != cudaSuccess) return cudaErrorUnknown; // old buffer may be in use + if (w.ws) cudaFree(w.ws); + if (w.counters) cudaFree(w.counters); + w.ws = nullptr; w.counters = nullptr; w.ws_elems = 0; w.n_counters = 0; + const size_t e = elems > w.ws_elems ? elems : w.ws_elems, c = counters > 4096 ? counters : 4096; + if (cudaMalloc(&w.ws, e * sizeof(int)) != cudaSuccess) return cudaErrorMemoryAllocation; + if (cudaMalloc(&w.counters, c * sizeof(int)) != cudaSuccess) return cudaErrorMemoryAllocation; + // Zeroed on the caller's stream: ggml's streams are non-blocking, so a + // legacy-stream memset would not be ordered before the first split launch. + if (cudaMemsetAsync(w.ws, 0, e * sizeof(int), stream) != cudaSuccess || + cudaMemsetAsync(w.counters, 0, c * sizeof(int), stream) != cudaSuccess) + return cudaErrorUnknown; + w.ws_elems = e; w.n_counters = c; + return cudaSuccess; +} + +template +cudaError_t launch_tile(const GemmArgs & g_in, cudaStream_t stream) { + using T = Tile; + // Opt into the dynamic shared memory once per instantiation. Not a stream + // operation, so it is safe under CUDA-graph capture. + static bool attr_set = false; + if (!attr_set) { + cudaError_t e = cudaFuncSetAttribute(gemm_i8_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, T::SMEM_BYTES); + if (e == cudaSuccess) + e = cudaFuncSetAttribute(gemm_i8_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, T::SMEM_BYTES); + if (e != cudaSuccess) return e; + attr_set = true; + } + GemmArgs g = g_in; + const dim3 grid((unsigned) ((g.M + BM - 1) / BM), (unsigned) (g.N / BN), (unsigned) SPLITK); + if (SPLITK > 1) { + if ((g.K / BK) % SPLITK != 0) return cudaErrorNotSupported; + const cudaError_t e = ensure_split_ws((size_t) g.M * (size_t) g.N, (size_t) grid.x * grid.y, stream); + if (e != cudaSuccess) return e; + g.ws = split_ws().ws; g.counters = split_ws().counters; + } + if (g.bias) gemm_i8_kernel<<>>(g); + else gemm_i8_kernel<<>>(g); + return cudaGetLastError(); +} + +} // namespace + +// VLA_FQ_TILE=: force one tiling for every shape (tuning aid; see the +// table in foldquant_gemm_check). Unset or -1 = the per-shape dispatch below. +static int forced_tile() { + static const int v = [] { + const char * e = std::getenv("VLA_FQ_TILE"); + return e && *e ? std::atoi(e) : -1; + }(); + return v; +} + +static cudaError_t launch_forced(int t, const GemmArgs & g, cudaStream_t stream) { + switch (t) { + case 0: return launch_tile<256, 64, 64, 3, 4, 2>(g, stream); + case 1: return launch_tile<64, 32, 128, 3, 2, 2>(g, stream); + case 2: return launch_tile<64, 64, 64, 4, 2, 2>(g, stream); + case 3: return launch_tile<128, 64, 64, 4, 4, 2>(g, stream); + case 4: return launch_tile<128, 64, 64, 3, 2, 2>(g, stream); + case 5: return launch_tile<128, 128, 64, 3, 4, 2>(g, stream); + case 6: return launch_tile<192, 64, 64, 3, 4, 2>(g, stream); + case 7: return launch_tile<64, 128, 64, 4, 2, 4>(g, stream); + case 8: return launch_tile<128, 64, 128, 2, 4, 2>(g, stream); + case 9: return launch_tile<64, 64, 128, 3, 2, 2>(g, stream); + case 10: return launch_tile<32, 64, 128, 4, 1, 2>(g, stream); + case 11: return launch_tile<64, 64, 64, 6, 2, 2>(g, stream); + case 12: return launch_tile<64, 32, 128, 6, 2, 2>(g, stream); + case 13: return launch_tile<64, 32, 256, 4, 2, 2>(g, stream); + case 14: return launch_tile<64, 64, 128, 6, 2, 2>(g, stream); + case 15: return launch_tile<64, 32, 128, 8, 2, 2>(g, stream); + case 16: return launch_tile<64, 32, 256, 6, 2, 2>(g, stream); + case 20: return launch_tile<64, 32, 128, 3, 2, 2, 2>(g, stream); + case 21: return launch_tile<64, 32, 128, 3, 2, 2, 3>(g, stream); + case 22: return launch_tile<64, 32, 128, 3, 2, 2, 4>(g, stream); + case 23: return launch_tile<64, 64, 64, 4, 2, 2, 2>(g, stream); + case 24: return launch_tile<64, 64, 64, 4, 2, 2, 4>(g, stream); + case 25: return launch_tile<64, 32, 128, 3, 2, 2, 6>(g, stream); + default: return cudaErrorInvalidValue; + } +} + +static int mma_variant() { + static const int v = [] { + const char * e = std::getenv("VLA_FQ_MMA"); + return e && *e ? std::atoi(e) : -1; + }(); + return v; +} +// VLA_FQ_MMA_SMALL_M=: mma variant for the M <= 64 shapes only (the DiT +// sites), leaving the prefill dispatch alone. +static int mma_variant_small_m() { + static const int v = [] { + const char * e = std::getenv("VLA_FQ_MMA_SMALL_M"); + return e && *e ? std::atoi(e) : -1; + }(); + return v; +} +static bool use_wmma() { + static const bool w = [] { const char * e = std::getenv("VLA_FQ_GEMM"); return e && std::string(e) == "wmma"; }(); + return w; +} + +cudaError_t launch_gemm(const GemmArgs & g, cudaStream_t stream) { + if (g.M <= 0) return cudaSuccess; + // The mma.sync kernel serves W8A8 and W4A4 (fq_gemm_mma.cu); the wmma + // tiles below stay as the VLA_FQ_GEMM=wmma fallback and for the tile sweep. + if (!use_wmma() && forced_tile() < 0) { + const int v = (g.M <= 64 && mma_variant_small_m() >= 0) ? mma_variant_small_m() : mma_variant(); + const cudaError_t e = launch_gemm_mma(g, v, stream); + if (e != cudaErrorNotSupported) return e; + } + if (g.wbits != 8 || g.abits != 8) return cudaErrorNotSupported; // W4A8 / mixed: CPU reference + if (g.K % 128 != 0 && g.K % 64 != 0) return cudaErrorNotSupported; + if (g.N % 64 != 0) return cudaErrorNotSupported; + // VLA_FQ_TILE_SMALL_M=1 restricts the override to the M <= 64 shapes (the + // DiT sites), so a model run times one DiT tiling with the LLM dispatch intact. + if (forced_tile() >= 0 && !(std::getenv("VLA_FQ_TILE_SMALL_M") && g.M > 64)) { + const int t = forced_tile(); + if (g.K % 128 != 0 && (t == 1 || t == 8 || t == 9 || t == 10 || t == 12 || t == 14 || t == 15)) return cudaErrorNotSupported; + if (g.K % 256 != 0 && (t == 13 || t == 16)) return cudaErrorNotSupported; + return launch_forced(t, g, stream); + } + // LLM prefill (M > 64, ~160 tokens on GR00T): measured on Orin (VLA_FQ_TILE + // sweep, tests/foldquant_gemm_check). The wide gate+up site (N = 12288) and + // long prefixes want the 128-row, BK=128 two-stage tile (M=160: 719 -> 426 + // us; M=1024 qkv 1081 -> 689 us); the other sites a 192-row tile that holds + // 160 tokens in one M tile with 17% zero rows instead of 256's 37% + // (qkv 173 -> 131, o 122 -> 69, down 339 -> 207 us). + if (g.M > 64) { + if ((g.N >= 8192 || g.M > 256) && g.K % 128 == 0) + return launch_tile<128, 64, 128, 2, 4, 2>(g, stream); + return launch_tile<192, 64, 64, 3, 4, 2>(g, stream); + } + // M <= 64 (the DiT's 41 action tokens): narrow N (o_proj, ff2, cross kv) + // gets a 32-wide N tile for the CTA count. Split-K over the same tiles + // (VLA_FQ_TILE=20..25, exact int32 atomics) was measured 5-8% slower end + // to end on Orin - the atomics and the extra L2 round trip cost more than + // the added CTAs buy - so it stays an opt-in candidate for other GPUs. + if (g.N <= 3072 && g.K % 128 == 0) + return launch_tile<64, 32, 128, 3, 2, 2>(g, stream); + return launch_tile<64, 64, 64, 4, 2, 2>(g, stream); +} + +} // namespace fq +} // namespace vla diff --git a/src/kernels/foldquant/fq_gemm_mma.cu b/src/kernels/foldquant/fq_gemm_mma.cu new file mode 100644 index 0000000..6d702b9 --- /dev/null +++ b/src/kernels/foldquant/fq_gemm_mma.cu @@ -0,0 +1,324 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Integer tensor-core GEMM on mma.sync + ldmatrix, the successor of the wmma +// kernel in fq_gemm_i8.cu for the shapes a VLA runs: a few dozen to a few +// hundred activation rows against a weight matrix streamed once from DRAM. +// +// Two operand widths share one kernel body: +// INT8: mma.m16n8k32.s8 - smem rows hold BK bytes (one byte per k) +// INT4: mma.m16n8k64.s4 - smem rows hold BK/2 bytes (two k per byte, the +// low nibble the even k, exactly the file and blob packing) +// and the same ldmatrix fragment mapping serves both, since ldmatrix moves +// 16-byte row segments and mma reads its A/B registers as packed k runs. +// +// Layout: A = the activation blob (row m at m*row_bytes, codes then the float +// scale), B = W[N][K_pack] with K contiguous. Both tiles land in shared memory +// through a cp.async pipeline with the 16-byte chunks XOR-swizzled by (row & 7), +// so the eight rows an ldmatrix touches hit eight different bank groups. +// Accumulation is int32 and the epilogue y = ((float)acc * act_scale[m]) * +// wscale[n] (+ bias[n]) is applied per element in that order, so the result is +// bit-identical to the wmma kernel and the CPU reference for any tiling. + +#include "fq_kernels.h" + +#include + +namespace vla { +namespace fq { + +namespace { + +__device__ __forceinline__ void cp_async16(void * smem, const void * gmem, bool pred) { + const unsigned s = (unsigned) __cvta_generic_to_shared(smem); + const int sz = pred ? 16 : 0; // src-size 0 zero-fills the 16 bytes + asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" :: "r"(s), "l"(gmem), "r"(sz)); +} +__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); } +template +__device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); } + +__device__ __forceinline__ void ldmatrix_x4(unsigned & r0, unsigned & r1, unsigned & r2, unsigned & r3, const void * smem) { + const unsigned s = (unsigned) __cvta_generic_to_shared(smem); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(s)); +} + +// D = A * B + D on one m16n8 tile; K is 32 (s8) or 64 (s4) packed values. +template +__device__ __forceinline__ void mma_tile(int * c, const unsigned * a, const unsigned * b) { + if (KPACK == 1) { + asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+r"(c[0]), "+r"(c[1]), "+r"(c[2]), "+r"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + } else { + asm volatile("mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+r"(c[0]), "+r"(c[1]), "+r"(c[2]), "+r"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + } +} + +// BM x BN CTA tile, BKB bytes of K per stage (32 or 64 k-values per mma step, +// i.e. 32 bytes for both widths), WM x WN warps. Rows in smem are BKB bytes, +// a multiple of 128 so the swizzle covers the 8 chunks a bank group spans. +template +struct MTile { + static constexpr int THREADS = WM * WN * 32; + static constexpr int CHUNKS_R = BKB / 16; // 16-byte chunks per row + static constexpr int STAGE_BYTES = (BM + BN) * BKB; + static constexpr int SMEM_BYTES = STAGES * STAGE_BYTES; + static constexpr int CHUNKS = (BM + BN) * CHUNKS_R; + static constexpr int PER_THR = CHUNKS / THREADS; + static constexpr int MT = BM / WM / 16; // m16 tiles per warp + static constexpr int NT = BN / WN / 8; // n8 tiles per warp + static constexpr int KSTEP_B = 32; // bytes of k per mma step + static constexpr int KSTEPS = BKB / KSTEP_B; + static_assert(CHUNKS % THREADS == 0, "tile copies must split evenly over the threads"); + static_assert(BM % (WM * 16) == 0 && BN % (WN * 16) == 0, "warp tile: multiples of m16 / n16"); + static_assert(BKB % 128 == 0, "smem rows must be a multiple of 128 bytes for the swizzle"); + static_assert(NT % 2 == 0, "B fragments load two n8 tiles per ldmatrix.x4"); +}; + +// Physical byte offset of 16-byte chunk c of row r within a tile whose rows are BKB bytes. +template +__device__ __forceinline__ int swz(int r, int c) { return r * BKB + ((c ^ (r & 7)) << 4); } + +template +__global__ void __launch_bounds__(WM * WN * 32) gemm_mma_kernel(const GemmArgs g) { + using T = MTile; + extern __shared__ __align__(128) int8_t smem[]; + + const int tid = threadIdx.x; + const int warp = tid >> 5, lane = tid & 31; + const int wm = warp / WN, wn = warp % WN; + const int64_t m0 = (int64_t) blockIdx.x * BM; + const int64_t n0 = (int64_t) blockIdx.y * BN; + const int64_t KB = g.K / KPACK; // packed bytes per row + const int KT = (int) (KB / BKB); + + // Per-thread copy list for one stage (sources at k0 = 0). + const int8_t * src[T::PER_THR]; + int dst[T::PER_THR]; + bool pred[T::PER_THR]; + #pragma unroll + for (int i = 0; i < T::PER_THR; ++i) { + const int c = tid + i * T::THREADS; + const int r = c / T::CHUNKS_R, c16 = c % T::CHUNKS_R; + if (r < BM) { + const int64_t m = m0 + r; + pred[i] = m < g.M; + src[i] = g.blob + (pred[i] ? m : 0) * g.row_bytes + c16 * 16; + dst[i] = swz(r, c16); + } else { + const int rb = r - BM; + pred[i] = true; + src[i] = g.w + (n0 + rb) * KB + c16 * 16; + dst[i] = BM * BKB + swz(rb, c16); + } + } + auto load_tile = [&](int stage, int kt) { + int8_t * base = smem + stage * T::STAGE_BYTES; + const int64_t k0 = (int64_t) kt * BKB; + #pragma unroll + for (int i = 0; i < T::PER_THR; ++i) cp_async16(base + dst[i], src[i] + k0, pred[i]); + }; + + int acc[T::MT][T::NT][4]; + #pragma unroll + for (int i = 0; i < T::MT; ++i) + #pragma unroll + for (int j = 0; j < T::NT; ++j) + #pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = 0; + + #pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < KT) load_tile(s, s); + cp_async_commit(); + } + + // ldmatrix row addressing: lane l provides the row (l & 7) of matrix (l >> 3). + // A matrices per m16 x 32B step: (rows 0-7, chunk 0), (rows 8-15, chunk 0), + // (rows 0-7, chunk 1), (rows 8-15, chunk 1) -> a0..a3. + // B matrices per n16 x 32B step: (n 0-7, chunk 0), (n 0-7, chunk 1), + // (n 8-15, chunk 0), (n 8-15, chunk 1) -> b0,b1 of tile n, b0,b1 of tile n+8. + const int a_row = (lane & 7) + ((lane >> 3) & 1) * 8, a_chk = lane >> 4; + const int b_row = (lane & 7) + (lane >> 4) * 8, b_chk = (lane >> 3) & 1; + + for (int kt = 0; kt < KT; ++kt) { + cp_async_wait(); + __syncthreads(); + const int stage = kt % STAGES; + const int nk = kt + STAGES - 1; + if (nk < KT) load_tile(nk % STAGES, nk); + cp_async_commit(); + + const int8_t * as = smem + stage * T::STAGE_BYTES; + const int8_t * bs = as + BM * BKB; + #pragma unroll + for (int ks = 0; ks < T::KSTEPS; ++ks) { + unsigned af[T::MT][4], bf[T::NT][2]; + #pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int r = wm * (T::MT * 16) + i * 16 + a_row; + ldmatrix_x4(af[i][0], af[i][1], af[i][2], af[i][3], as + swz(r, ks * 2 + a_chk)); + } + #pragma unroll + for (int j = 0; j < T::NT; j += 2) { + const int r = wn * (T::NT * 8) + j * 8 + b_row; + ldmatrix_x4(bf[j][0], bf[j][1], bf[j + 1][0], bf[j + 1][1], bs + swz(r, ks * 2 + b_chk)); + } + #pragma unroll + for (int i = 0; i < T::MT; ++i) + #pragma unroll + for (int j = 0; j < T::NT; ++j) + mma_tile(acc[i][j], af[i], bf[j]); + } + } + cp_async_wait<0>(); + + // Next site's weights: every CTA prefetches an equal slice of the head of + // that matrix into L2 now that its own loads are in, so the next GEMM's + // first stages hit L2 instead of DRAM. Fire-and-forget; no data dependence. + if (g.pf) { + const int64_t lines = g.pf_bytes >> 7; + const int64_t ctas = (int64_t) gridDim.x * gridDim.y; + const int64_t cta = (int64_t) blockIdx.y * gridDim.x + blockIdx.x; + const int64_t per = (lines + ctas - 1) / ctas; + const int64_t l0 = cta * per, l1 = l0 + per < lines ? l0 + per : lines; + for (int64_t l = l0 + tid; l < l1; l += T::THREADS) + asm volatile("prefetch.global.L2 [%0];" :: "l"(g.pf + (l << 7))); + } + + // Fragment layout of a m16n8 int32 tile: c0,c1 at (row g, cols 2q, 2q+1), + // c2,c3 at (row g+8, same cols), g = lane/4, q = lane%4. + const int grp = lane >> 2, q = lane & 3; + // The per-column factors are shared by every row this thread writes: load + // them once per n8 tile instead of once per (row, tile). + float2 ws[T::NT], bs2[T::NT]; + int ho[T::NT], hs[T::NT]; // head layout: y[ho + m*hs] for column n, column n+1 at +1 (Q/K) or +M (V) + #pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int64_t n = n0 + wn * (T::NT * 8) + j * 8 + 2 * q; + ws[j] = *(const float2 *) (g.wscale + n); + if (HAS_BIAS) bs2[j] = *(const float2 *) (g.bias + n); + if (g.heads) fq_out_column(g, (int) n, ho[j], hs[j]); + } + #pragma unroll + for (int i = 0; i < T::MT; ++i) { + #pragma unroll + for (int h = 0; h < 2; ++h) { + const int64_t m = m0 + wm * (T::MT * 16) + i * 16 + grp + h * 8; + if (m >= g.M) continue; + const float xs = *(const float *) (g.blob + m * g.row_bytes + KB); + #pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int64_t n = n0 + wn * (T::NT * 8) + j * 8 + 2 * q; + float v0 = ((float) acc[i][j][2 * h] * xs) * ws[j].x; + float v1 = ((float) acc[i][j][2 * h + 1] * xs) * ws[j].y; + if (HAS_BIAS) { v0 = v0 + bs2[j].x; v1 = v1 + bs2[j].y; } + if (g.res) { // the residual add that followed the GEMM: one float add, same result + const float2 r = *(const float2 *) (g.res + m * g.N + n); + v0 = v0 + r.x; v1 = v1 + r.y; + } + if (g.heads) { // head layout (hd even, n even: a Q/K pair stays adjacent and 8-B aligned) + float * yp = g.y + ho[j] + m * hs[j]; + if (hs[j] == 1) { yp[0] = v0; yp[(int) g.M] = v1; } + else { *(float2 *) yp = make_float2(v0, v1); } + } else { + *(float2 *) (g.y + m * g.N + n) = make_float2(v0, v1); + } + } + } + } +} + +template +cudaError_t launch_mma(const GemmArgs & g, cudaStream_t stream) { + using T = MTile; + static bool attr_set = false; + if (!attr_set) { + cudaError_t e = cudaFuncSetAttribute(gemm_mma_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, T::SMEM_BYTES); + if (e == cudaSuccess) + e = cudaFuncSetAttribute(gemm_mma_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, T::SMEM_BYTES); + if (e != cudaSuccess) return e; + attr_set = true; + } + if ((g.K / KPACK) % BKB != 0 || g.N % BN != 0) return cudaErrorNotSupported; + const dim3 grid((unsigned) ((g.M + BM - 1) / BM), (unsigned) (g.N / BN)); + if (g.bias) gemm_mma_kernel<<>>(g); + else gemm_mma_kernel<<>>(g); + return cudaGetLastError(); +} + +} // namespace + +// variant: -1 = the per-shape default, otherwise one of the candidates below +// (VLA_FQ_MMA= from launch_gemm), for tuning on another GPU. +cudaError_t launch_gemm_mma(const GemmArgs & g, int variant, cudaStream_t stream) { + const bool w4 = g.wbits == 4; + if ((g.wbits != 8 && g.wbits != 4) || g.abits != g.wbits) return cudaErrorNotSupported; // W8A8 or W4A4 + if (variant < 0) { + if (g.M > 64) { + // 128x64x128, 3 stages (72 KB smem, 2 CTAs/SM) for every prefill site: on + // Orin it beat the 192-row tile end to end (ALOHA shape p50 174 vs 205 ms). + variant = 3; + } else { + variant = g.N <= 3072 ? 0 : 1; + } + } + if (!w4) { + switch (variant) { + case 0: return launch_mma<64, 32, 128, 3, 2, 2, 1>(g, stream); + case 1: return launch_mma<64, 64, 128, 3, 2, 2, 1>(g, stream); + case 2: return launch_mma<192, 64, 128, 3, 4, 2, 1>(g, stream); + case 3: return launch_mma<128, 64, 128, 3, 4, 2, 1>(g, stream); + case 4: return launch_mma<64, 32, 256, 3, 2, 2, 1>(g, stream); + case 5: return launch_mma<128, 32, 128, 3, 4, 2, 1>(g, stream); + case 6: return launch_mma<64, 64, 128, 4, 2, 2, 1>(g, stream); + case 7: return launch_mma<256, 64, 128, 3, 4, 2, 1>(g, stream); + case 8: return launch_mma<192, 64, 128, 2, 4, 2, 1>(g, stream); // 64 KB smem: 2 CTAs/SM + case 9: return launch_mma<128, 64, 128, 2, 4, 2, 1>(g, stream); + case 10: return launch_mma<96, 64, 128, 3, 2, 2, 1>(g, stream); // 3 m16 tiles per warp + // Small-M candidates (DiT, M = 41): fewer padded rows, more CTAs, deeper + // pipeline. Cold micro-bench on Orin: 11/13 within 2% of the default, the + // rest 5-45% slower (the sites run at 130-145 GB/s, near the DRAM ceiling). + case 11: return launch_mma<48, 32, 128, 3, 1, 2, 1>(g, stream); // 2 warps, 30 KB smem + case 12: return launch_mma<48, 32, 128, 6, 1, 2, 1>(g, stream); + case 13: return launch_mma<48, 64, 128, 3, 1, 4, 1>(g, stream); + case 14: return launch_mma<32, 32, 128, 3, 1, 2, 1>(g, stream); // two m-tiles per site + case 15: return launch_mma<64, 32, 128, 6, 2, 2, 1>(g, stream); + case 16: return launch_mma<48, 32, 256, 3, 1, 2, 1>(g, stream); + default: return cudaErrorInvalidValue; + } + } + switch (variant) { // INT4: rows hold K/2 bytes, so BKB=128 covers 256 k + case 0: return launch_mma<64, 32, 128, 3, 2, 2, 2>(g, stream); + case 1: return launch_mma<64, 64, 128, 3, 2, 2, 2>(g, stream); + case 2: return launch_mma<192, 64, 128, 3, 4, 2, 2>(g, stream); + case 11: return launch_mma<48, 32, 128, 3, 1, 2, 2>(g, stream); + case 12: return launch_mma<48, 32, 128, 6, 1, 2, 2>(g, stream); + case 13: return launch_mma<48, 64, 128, 3, 1, 4, 2>(g, stream); + case 14: return launch_mma<32, 32, 128, 3, 1, 2, 2>(g, stream); + case 15: return launch_mma<64, 32, 128, 6, 2, 2, 2>(g, stream); + case 16: return launch_mma<48, 32, 256, 3, 1, 2, 2>(g, stream); + case 3: return launch_mma<128, 64, 128, 3, 4, 2, 2>(g, stream); + default: return launch_mma<64, 32, 128, 3, 2, 2, 2>(g, stream); + } +} + +} // namespace fq +} // namespace vla diff --git a/src/kernels/foldquant/fq_kernels.h b/src/kernels/foldquant/fq_kernels.h new file mode 100644 index 0000000..26ede73 --- /dev/null +++ b/src/kernels/foldquant/fq_kernels.h @@ -0,0 +1,97 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant CUDA kernels: the activation prologue ([RMSNorm] -> [/s] -> block +// FWHT -> [/s] -> per-token INT8) and the INT8 x INT8 -> INT32 GEMM with the +// dequant epilogue. Pure CUDA, no ggml: the glue that decodes graph nodes lives +// in src/cuda/vla_cuda_foldquant.cu. +// +// Numerics are pinned to src/foldquant_ref.h: this archive is compiled with +// -fmad=false and without --use_fast_math so every float op rounds exactly +// where the reference does. All launches run on the caller's stream, allocate +// nothing and never synchronize (CUDA-graph capture safe). + +#pragma once + +#include + +#include + +namespace vla { +namespace fq { + +struct ActArgs { + const float * x; // [M][K] F32; rows contiguous, x_stride floats apart (0 = K) + int64_t x_stride = 0; + const float * ascale; // [K] or null + const float * gamma; // [K] or null (fused RMSNorm) + int8_t * blob; // [M][row_bytes]: codes then the float scale at K_pack + int64_t row_bytes; + int64_t M, K; + int abits; // 8 | 4 + int rot_block; // 1 = none + bool fold_before; + float clip, eps; + float inv_sqrt_bs; +}; + +struct GemmArgs { + const int8_t * w; // [N][K_pack] INT8 codes (nibbles when wbits == 4) + const int8_t * blob; // activation blob from the prologue + const float * wscale; // [N] + const float * bias; // [N] or null + float * y; // [M][N] F32 + const float * res = nullptr; // [M][N] F32 residual added in the epilogue (y = ... + res), or null + const int8_t * pf = nullptr; // next site's weights to prefetch into L2 (pf_bytes of them), or null + int64_t pf_bytes = 0; + // Head layout of y (0 = plain [M][N]); see FqGemmSpec / fq_out_index. + int head_dim = 0, heads = 0; + uint32_t vmask = 0; + int64_t M, N, K; + int64_t row_bytes; + int wbits, abits; + // Split-K only (set by launch_gemm): int32 [M][N] partial-sum workspace and + // one arrival counter per output tile, both zero between launches. + int * ws = nullptr; + int * counters = nullptr; +}; + +#ifdef __CUDACC__ +// Column n of a head-laid-out output maps to y[off + m * stride]: this splits +// the column part so the epilogue can hoist it out of its row loop. 32-bit +// math: M * N < 2^31 for every site. +__device__ __forceinline__ void fq_out_column(const GemmArgs & g, int n, int & off, int & stride) { + const int M = (int) g.M, dim = g.head_dim * g.heads; + const int p = n / dim, r = n - p * dim, h = r / g.head_dim, d = r - h * g.head_dim; + const int base = p * dim * M; + if ((g.vmask >> p) & 1) { off = base + (h * g.head_dim + d) * M; stride = 1; } + else { off = base + h * M * g.head_dim + d; stride = g.head_dim; } +} +__device__ __forceinline__ int64_t fq_out_index(const GemmArgs & g, int64_t m, int64_t n) { + if (!g.heads) return m * g.N + n; + int off, stride; + fq_out_column(g, (int) n, off, stride); + return (int64_t) off + m * (int64_t) stride; +} +#endif + +// Both return cudaSuccess or the launch error. Shapes the kernels do not cover +// (W4/A4 until phase 2/3) return cudaErrorNotSupported without launching. +cudaError_t launch_act (const ActArgs & a, cudaStream_t stream); +cudaError_t launch_gemm(const GemmArgs & g, cudaStream_t stream); +// mma.sync/ldmatrix kernel (fq_gemm_mma.cu); W8A8 and W4A4. variant -1 = default. +cudaError_t launch_gemm_mma(const GemmArgs & g, int variant, cudaStream_t stream); + +} // namespace fq +} // namespace vla diff --git a/src/kernels/foldquant/fq_prologue.cu b/src/kernels/foldquant/fq_prologue.cu new file mode 100644 index 0000000..3bf7ed2 --- /dev/null +++ b/src/kernels/foldquant/fq_prologue.cu @@ -0,0 +1,357 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// One warp per token row, four rows per CTA. Lane l owns the 64-element chunks +// c = l, l+32, ... of the row, so a chunk's butterfly (block 64 or smaller) +// runs entirely in that lane's registers: no shared memory, no barriers, and +// the two row-wide reductions (sum of squares, amax) are warp shuffles. The +// reduction tree - per-lane partial in element order, xor butterfly 16..1 - +// is the one foldquant_ref.h mirrors, so rstd is bit-identical between +// backends. amax is exact in any order. +// +// The row is read twice (three times with a fused RMSNorm): once per pass the +// value must be known for. It is L1/L2 resident, so the passes cost ALU, not +// bandwidth, which is what matters at M of a few hundred rows. + +#include "fq_kernels.h" + +#include + +namespace vla { +namespace fq { + +namespace { + +constexpr int CHUNK = 64; +constexpr int WARPS = 4; +constexpr int THREADS = WARPS * 32; + +__device__ __forceinline__ float warp_sum(float v) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + v = v + __shfl_xor_sync(0xffffffffu, v, off); + return v; +} + +__device__ __forceinline__ float warp_max(float v) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off)); + return v; +} + +// Load a 64-float chunk, apply the pre-quant chain, return it rotated. ROT is +// a compile-time constant so every index into v[] is static and the chunk +// stays in registers (a runtime loop bound would spill it to local memory). +template +__device__ __forceinline__ void load_chunk(const ActArgs & a, const float * x, int64_t c0, float rstd, float * v) { + const float4 * src = (const float4 *) (x + c0); + #pragma unroll + for (int i = 0; i < CHUNK / 4; ++i) { + const float4 q = src[i]; + v[4*i] = q.x; v[4*i+1] = q.y; v[4*i+2] = q.z; v[4*i+3] = q.w; + } + if (a.gamma) { + const float4 * g = (const float4 *) (a.gamma + c0); + #pragma unroll + for (int i = 0; i < CHUNK / 4; ++i) { + const float4 q = g[i]; + v[4*i] = (v[4*i] * rstd) * q.x; + v[4*i+1] = (v[4*i+1] * rstd) * q.y; + v[4*i+2] = (v[4*i+2] * rstd) * q.z; + v[4*i+3] = (v[4*i+3] * rstd) * q.w; + } + } + if (a.ascale && a.fold_before) { + const float4 * s = (const float4 *) (a.ascale + c0); + #pragma unroll + for (int i = 0; i < CHUNK / 4; ++i) { + const float4 q = s[i]; + v[4*i] = v[4*i] / q.x; v[4*i+1] = v[4*i+1] / q.y; v[4*i+2] = v[4*i+2] / q.z; v[4*i+3] = v[4*i+3] / q.w; + } + } + if (ROT > 1) { + // Natural-order Sylvester butterfly on every ROT sub-block of the chunk. + #pragma unroll + for (int h = 1; h < ROT; h <<= 1) { + #pragma unroll + for (int p = 0; p < CHUNK / 2; ++p) { + // pairs (j, j+h) with j = (p / h) * 2h + (p % h) + const int j = ((p / h) * 2 * h) + (p % h); + const float u = v[j], w = v[j + h]; + v[j] = u + w; + v[j + h] = u - w; + } + } + #pragma unroll + for (int i = 0; i < CHUNK; ++i) v[i] = v[i] * a.inv_sqrt_bs; + } + if (a.ascale && !a.fold_before) { + const float4 * s = (const float4 *) (a.ascale + c0); + #pragma unroll + for (int i = 0; i < CHUNK / 4; ++i) { + const float4 q = s[i]; + v[4*i] = v[4*i] / q.x; v[4*i+1] = v[4*i+1] / q.y; v[4*i+2] = v[4*i+2] / q.z; v[4*i+3] = v[4*i+3] / q.w; + } + } +} + +template +__global__ void __launch_bounds__(THREADS) act_kernel(const ActArgs a) { + const int lane = threadIdx.x & 31; + const int64_t m = (int64_t) blockIdx.x * WARPS + (threadIdx.x >> 5); + if (m >= a.M) return; + const int64_t K = a.K; + const int nchunks = (int) (K / CHUNK); + const float * x = a.x + m * (a.x_stride ? a.x_stride : K); + + float rstd = 0.0f; + if (a.gamma) { + float p = 0.0f; + for (int c = lane; c < nchunks; c += 32) { + const float4 * src = (const float4 *) (x + (int64_t) c * CHUNK); + #pragma unroll + for (int i = 0; i < CHUNK / 4; ++i) { + const float4 q = src[i]; + p = p + q.x * q.x; p = p + q.y * q.y; p = p + q.z * q.z; p = p + q.w * q.w; + } + } + const float sumsq = warp_sum(p); + rstd = 1.0f / sqrtf(sumsq / (float) K + a.eps); + } + + float v[CHUNK]; + float mx = 0.0f; + for (int c = lane; c < nchunks; c += 32) { + load_chunk(a, x, (int64_t) c * CHUNK, rstd, v); + #pragma unroll + for (int i = 0; i < CHUNK; ++i) mx = fmaxf(mx, fabsf(v[i])); + } + mx = warp_max(mx); + + const float qmax = ABITS == 4 ? 7.0f : 127.0f; + float scale = (a.clip * mx) / qmax; + if (scale < 1e-12f) scale = 1e-12f; + // Quantize with the reciprocal, as FoldQuant's TensorRT kernels do + // (rmsnorm_per_row_quant_cuda.cu) and as foldquant_ref.h mirrors. + const float inv = 1.0f / scale; + + int8_t * row = a.blob + m * a.row_bytes; + for (int c = lane; c < nchunks; c += 32) { + // A lane that owns a single chunk still holds it rotated in v[]. + if (nchunks > 32) load_chunk(a, x, (int64_t) c * CHUNK, rstd, v); + if (ABITS == 8) { + uint32_t * dst = (uint32_t *) (row + (int64_t) c * CHUNK); + #pragma unroll + for (int i = 0; i < CHUNK; i += 4) { + uint32_t packed = 0; + #pragma unroll + for (int j = 0; j < 4; ++j) { + float q = rintf(v[i + j] * inv); + q = fminf(qmax, fmaxf(-qmax, q)); + packed |= ((uint32_t) (int) q & 0xFFu) << (8 * j); + } + dst[i / 4] = packed; + } + } else { + uint32_t * dst = (uint32_t *) (row + (int64_t) c * (CHUNK / 2)); + #pragma unroll + for (int i = 0; i < CHUNK; i += 8) { + uint32_t packed = 0; + #pragma unroll + for (int j = 0; j < 8; j += 2) { + float q0 = rintf(v[i + j] * inv), q1 = rintf(v[i + j + 1] * inv); + q0 = fminf(qmax, fmaxf(-qmax, q0)); + q1 = fminf(qmax, fmaxf(-qmax, q1)); + const uint32_t b = ((uint32_t) (int) q0 & 0xFu) | (((uint32_t) (int) q1 & 0xFu) << 4); + packed |= b << (4 * j); + } + dst[i / 8] = packed; + } + } + } + if (lane == 0) *(float *) (row + (ABITS == 8 ? K : K / 2)) = scale; +} + +// CTA-per-row variant for the short rows a VLA quantizes (41 action tokens, +// a few hundred LLM tokens): one warp per row leaves a 16-SM GPU nearly idle +// and each launch latency-bound at ~50 us. Here RW warps share a row; a warp +// takes chunks c = warp, warp + RW, ... and its lanes own the element pair +// (2*lane, 2*lane + 1) of each chunk, so loads and stores are coalesced and +// the per-row chain is a handful of shuffles. The 64-wide butterfly runs +// stage h = 1 inside the lane and the other stages across lanes: every stage +// forms exactly the (a + b, a - b) pairs fwht_row forms, so the values are +// bit-identical to the reference. The RMSNorm sum of squares keeps the +// reference's lane-partial order (warp 0 runs the warp-per-row loop). +constexpr int RW = 8; // warps per row +constexpr int RT = RW * 32; +constexpr int MAXC = 16; // chunks per warp held in registers: K <= RW * MAXC * 64 + +template +__device__ __forceinline__ void rot_pair(float & v0, float & v1, int lane, float inv_sqrt_bs) { + if (ROT > 1) { + { const float u = v0, w = v1; v0 = u + w; v1 = u - w; } // h = 1: pair (2l, 2l+1) + #pragma unroll + for (int h = 2; h < ROT; h <<= 1) { + const int half = h >> 1; // partner lane distance + const float o0 = __shfl_xor_sync(0xffffffffu, v0, half); + const float o1 = __shfl_xor_sync(0xffffffffu, v1, half); + const bool upper = (lane & half) != 0; // this lane holds j + h + v0 = upper ? (o0 - v0) : (v0 + o0); + v1 = upper ? (o1 - v1) : (v1 + o1); + } + v0 = v0 * inv_sqrt_bs; + v1 = v1 * inv_sqrt_bs; + } +} + +template +__global__ void __launch_bounds__(RT) act_row_kernel(const ActArgs a) { + __shared__ float red[RW]; + __shared__ float s_rstd; + const int tid = threadIdx.x; + const int warp = tid >> 5, lane = tid & 31; + const int64_t m = blockIdx.x; + const int64_t K = a.K; + const int nchunks = (int) (K / CHUNK); + const float * x = a.x + m * (a.x_stride ? a.x_stride : K); + + if (a.gamma) { + if (warp == 0) { + float p = 0.0f; + for (int c = lane; c < nchunks; c += 32) { + const float4 * src = (const float4 *) (x + (int64_t) c * CHUNK); + #pragma unroll + for (int i = 0; i < CHUNK / 4; ++i) { + const float4 q = src[i]; + p = p + q.x * q.x; p = p + q.y * q.y; p = p + q.z * q.z; p = p + q.w * q.w; + } + } + const float sumsq = warp_sum(p); + if (lane == 0) s_rstd = 1.0f / sqrtf(sumsq / (float) K + a.eps); + } + __syncthreads(); + } + const float rstd = a.gamma ? s_rstd : 0.0f; + + float v0[MAXC], v1[MAXC]; + float mx = 0.0f; + #pragma unroll + for (int i = 0; i < MAXC; ++i) { + const int c = warp + i * RW; + if (c < nchunks) { + const int64_t k0 = (int64_t) c * CHUNK + 2 * lane; + const float2 q = *(const float2 *) (x + k0); + float e0 = q.x, e1 = q.y; + if (a.gamma) { + const float2 g = *(const float2 *) (a.gamma + k0); + e0 = (e0 * rstd) * g.x; + e1 = (e1 * rstd) * g.y; + } + if (a.ascale && a.fold_before) { + const float2 s = *(const float2 *) (a.ascale + k0); + e0 = e0 / s.x; e1 = e1 / s.y; + } + rot_pair(e0, e1, lane, a.inv_sqrt_bs); + if (a.ascale && !a.fold_before) { + const float2 s = *(const float2 *) (a.ascale + k0); + e0 = e0 / s.x; e1 = e1 / s.y; + } + v0[i] = e0; v1[i] = e1; + mx = fmaxf(mx, fmaxf(fabsf(e0), fabsf(e1))); + } + } + mx = warp_max(mx); + if (lane == 0) red[warp] = mx; + __syncthreads(); + mx = red[0]; + #pragma unroll + for (int w = 1; w < RW; ++w) mx = fmaxf(mx, red[w]); + + const float qmax = ABITS == 4 ? 7.0f : 127.0f; + float scale = (a.clip * mx) / qmax; + if (scale < 1e-12f) scale = 1e-12f; + const float inv = 1.0f / scale; + int8_t * row = a.blob + m * a.row_bytes; + #pragma unroll + for (int i = 0; i < MAXC; ++i) { + const int c = warp + i * RW; + if (c < nchunks) { + float q0 = rintf(v0[i] * inv), q1 = rintf(v1[i] * inv); + q0 = fminf(qmax, fmaxf(-qmax, q0)); + q1 = fminf(qmax, fmaxf(-qmax, q1)); + if (ABITS == 8) { + const uint16_t packed = (uint16_t) (((uint32_t) (int) q0 & 0xFFu) | (((uint32_t) (int) q1 & 0xFFu) << 8)); + *(uint16_t *) (row + (int64_t) c * CHUNK + 2 * lane) = packed; + } else { + row[(int64_t) c * (CHUNK / 2) + lane] = (int8_t) (((uint32_t) (int) q0 & 0xFu) | (((uint32_t) (int) q1 & 0xFu) << 4)); + } + } + } + if (tid == 0) *(float *) (row + (ABITS == 8 ? K : K / 2)) = scale; +} + +} // namespace + +template +static cudaError_t launch_rot_row(const ActArgs & a, cudaStream_t stream) { + const unsigned grid = (unsigned) a.M; + switch (a.rot_block) { + case 64: act_row_kernel<<>>(a); break; + case 32: act_row_kernel<<>>(a); break; + case 16: act_row_kernel<<>>(a); break; + case 8: act_row_kernel<<>>(a); break; + case 4: act_row_kernel<<>>(a); break; + case 2: act_row_kernel<<>>(a); break; + case 1: + case 0: act_row_kernel<<>>(a); break; + default: return cudaErrorNotSupported; + } + return cudaGetLastError(); +} + +template +static cudaError_t launch_rot(const ActArgs & a, unsigned grid, cudaStream_t stream) { + switch (a.rot_block) { + case 64: act_kernel<<>>(a); break; + case 32: act_kernel<<>>(a); break; + case 16: act_kernel<<>>(a); break; + case 8: act_kernel<<>>(a); break; + case 4: act_kernel<<>>(a); break; + case 2: act_kernel<<>>(a); break; + case 1: + case 0: act_kernel<<>>(a); break; + default: return cudaErrorNotSupported; + } + return cudaGetLastError(); +} + +cudaError_t launch_act(const ActArgs & a, cudaStream_t stream) { + if (a.M <= 0) return cudaSuccess; + if (a.K % CHUNK != 0 || a.rot_block > CHUNK) return cudaErrorNotSupported; + if (a.x_stride && (a.x_stride % 4 != 0)) return cudaErrorNotSupported; // float4 loads need 16-byte rows + // VLA_FQ_ACT_WARP=1 keeps the warp-per-row kernel (A/B and fallback). + static const bool warp_only = [] { const char * e = std::getenv("VLA_FQ_ACT_WARP"); return e && *e && *e != '0'; }(); + if (!warp_only && a.K / CHUNK <= RW * MAXC) { + if (a.abits == 8) return launch_rot_row<8>(a, stream); + if (a.abits == 4) return launch_rot_row<4>(a, stream); + } + const unsigned grid = (unsigned) ((a.M + WARPS - 1) / WARPS); + if (a.abits == 8) return launch_rot<8>(a, grid, stream); + if (a.abits == 4) return launch_rot<4>(a, grid, stream); + return cudaErrorNotSupported; +} + +} // namespace fq +} // namespace vla diff --git a/src/layers/fq_linear.h b/src/layers/fq_linear.h new file mode 100644 index 0000000..8200d21 --- /dev/null +++ b/src/layers/fq_linear.h @@ -0,0 +1,95 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant linears as graph nodes. Two GGML_OP_CUSTOM nodes per site so the +// quantized activation is an ordinary gallocr intermediate (shared by q/k/v +// or gate/up, no backend workspace, capture-safe), and one encoding serves both +// backends: the CPU backend runs the custom function, the CUDA backend claims +// the node by the magic word in its userdata. + +#pragma once + +#include "foldquant.h" +#include "foldquant_ref.h" + +#include "ggml.h" + +#include + +namespace vla { + +// x: F32 [K, T, ...] -> I8 blob [row_bytes, T*...]. With s.gamma set, x is the +// pre-norm hidden state and the RMSNorm is fused into this node. +inline ggml_tensor * fq_act(ggml_context * C, const FqLinear & s, ggml_tensor * x) { + GGML_ASSERT(x->type == GGML_TYPE_F32 && x->ne[0] == s.act.K); + // Rows contiguous and 16-byte aligned with the higher dims packed is enough + // (the prologue takes a row stride); anything else is made contiguous. + const bool rows_ok = x->nb[0] == sizeof(float) && x->nb[1] % 16 == 0 && + x->nb[2] == x->nb[1] * (size_t) x->ne[1] && x->nb[3] == x->nb[2] * (size_t) x->ne[2]; + ggml_tensor * xin = rows_ok ? x : ggml_cont(C, x); + const int64_t rows = ggml_nelements(xin) / s.act.K; + + // Sources are packed without holes: src[1] is the gamma when the norm is + // fused, then the ascale when the site ships one (fq_act_srcs() decodes). + ggml_tensor * args[3] = { xin, nullptr, nullptr }; + int n_args = 1; + if (s.gamma) args[n_args++] = s.gamma; + if (s.ascale) args[n_args++] = s.ascale; + ggml_tensor * t = ggml_custom_4d(C, GGML_TYPE_I8, fq_act_row_bytes(s.act.K, s.act.abits), rows, 1, 1, + args, n_args, fq_act_cpu, GGML_N_TASKS_MAX, (void *) &s.act); + ggml_format_name(t, "%s.fq_act", ggml_get_name(s.w)); + return t; +} + +// xq: blob from fq_act (any site sharing the same input transform) -> F32 [N, T]. +// residual: an F32 tensor shaped like the output that the model would add +// right after the GEMM; folded into the epilogue (one float add, so the +// result is the same as ggml_add would produce). VLA_FQ_NO_FUSE=1 keeps the +// separate add for A/B runs. +inline bool fq_fuse_residual() { + static const bool off = [] { const char * e = std::getenv("VLA_FQ_NO_FUSE"); return e && *e && *e != '0'; }(); + return !off; +} + +inline ggml_tensor * fq_gemm(ggml_context * C, const FqLinear & s, ggml_tensor * xq, ggml_tensor * residual = nullptr) { + GGML_ASSERT(xq->type == GGML_TYPE_I8 && xq->ne[0] == fq_act_row_bytes(s.act.K, s.act.abits)); + const int64_t T = xq->ne[1]; + const bool fuse = residual && fq_fuse_residual() && !s.gemm.heads && residual->type == GGML_TYPE_F32 && + ggml_is_contiguous(residual) && residual->ne[0] == s.gemm.N && ggml_nelements(residual) == s.gemm.N * T; + ggml_tensor * args[5] = { s.w, xq, s.wscale, s.bias, fuse ? residual : nullptr }; + ggml_tensor * y = ggml_custom_4d(C, GGML_TYPE_F32, s.gemm.N, T, 1, 1, + args, fuse ? 5 : 4, fq_gemm_cpu, GGML_N_TASKS_MAX, (void *) &s.gemm); + ggml_format_name(y, "%s.fq_gemm", ggml_get_name(s.w)); + return (residual && !fuse) ? ggml_add(C, residual, y) : y; +} + +// Part `part` of a head-laid-out GEMM output (fq_set_heads), as the tensor the +// attention takes: [hd, T, heads] for Q/K, [T, hd, heads] for a V part. +inline ggml_tensor * fq_head_view(ggml_context * C, ggml_tensor * y, const FqLinear & s, int part, int64_t T) { + GGML_ASSERT(s.gemm.heads && y->ne[1] == T); + const int64_t hd = s.gemm.head_dim, heads = s.gemm.heads; + const size_t off = (size_t) part * heads * hd * T * sizeof(float); + ggml_tensor * v = ((s.gemm.vmask >> part) & 1) + ? ggml_view_3d(C, y, T, hd, heads, T * sizeof(float), T * hd * sizeof(float), off) + : ggml_view_3d(C, y, hd, T, heads, hd * sizeof(float), hd * T * sizeof(float), off); + ggml_format_name(v, "%s.part%d", ggml_get_name(y), part); + return v; +} + +// y = W x (+ bias) (+ residual) +inline ggml_tensor * fq_linear(ggml_context * C, const FqLinear & s, ggml_tensor * x, ggml_tensor * residual = nullptr) { + return fq_gemm(C, s, fq_act(C, s, x), residual); +} + +} // namespace vla diff --git a/src/loader.cpp b/src/loader.cpp index 340fb2c..b9cbb58 100644 --- a/src/loader.cpp +++ b/src/loader.cpp @@ -26,6 +26,33 @@ namespace { constexpr size_t NAME_CAP = 256; +// Resident type of a float rebuild: the requested GEMM type when it is a float +// type, F32 otherwise. +ggml_type float_type(ggml_type want) { + return (want == GGML_TYPE_F32 || want == GGML_TYPE_F16 || want == GGML_TYPE_BF16) ? want : GGML_TYPE_F32; +} + +} + +void WeightLoader::as_float(const std::string & name, int64_t K, int64_t N, FloatMaker make) { + as_float_[name] = AsFloat{K, N, std::move(make)}; +} + +bool WeightLoader::read_resident(const std::string & name, ggml_type type, bool gemma_norm, std::vector & out) { + const auto it = as_float_.find(name); + if (it == as_float_.end()) { + out = g_.read_convert(name.c_str(), type, gemma_norm); + return !out.empty(); + } + const AsFloat & a = it->second; + std::vector f; + if (!a.make(f) || (int64_t) f.size() != a.K * a.N) { + std::fprintf(stderr, "vla(%s): float rebuild of %s failed\n", arch_, name.c_str()); + return false; + } + out.assign(ggml_row_size(type, a.K) * (size_t) a.N, 0); + ggml_quantize_chunk(type, f.data(), out.data(), 0, a.N, a.K, nullptr); + return true; } ggml_tensor * WeightLoader::declare(ggml_type want, bool required, bool gemma_norm, @@ -47,6 +74,28 @@ ggml_tensor * WeightLoader::declare(ggml_type want, bool required, bool gemma_no return nullptr; } + // INT8 in the file is a FoldQuant site (docs/QUANTIZATION.md), which only + // typed(GGML_TYPE_I8) may declare; a float GEMM declare would otherwise fail + // later in read_convert with a size mismatch that says nothing useful. + if (src->type == GGML_TYPE_I8 && want != GGML_TYPE_I8) { + // ... unless the site was registered to be read back as a float weight. + const auto it = as_float_.find(name); + if (it != as_float_.end()) { + ggml_tensor * t = ggml_new_tensor_2d(ctx_, float_type(want), it->second.K, it->second.N); + if (!t) { + std::fprintf(stderr, "vla(%s): ggml_new_tensor failed for %s\n", arch_, name); + ok_ = false; + return nullptr; + } + ggml_set_name(t, name); + return t; + } + std::fprintf(stderr, "vla(%s): %s is INT8 (FoldQuant) but this site is not FoldQuant-aware in this arch\n", + arch_, name); + ok_ = false; + return nullptr; + } + ggml_tensor * t = ggml_new_tensor(ctx_, g_.resident_type(src, want), ggml_n_dims(src), src->ne); if (!t) { std::fprintf(stderr, "vla(%s): ggml_new_tensor failed for %s (weight context too small?)\n", arch_, name); @@ -85,6 +134,23 @@ ggml_tensor * WeightLoader::typed(ggml_type want, const char * fmt, ...) { return t; } +ggml_tensor * WeightLoader::opt_typed(ggml_type want, const char * fmt, ...) { + va_list ap; + va_start(ap, fmt); + ggml_tensor * t = declare(want, false, false, fmt, ap); + va_end(ap); + return t; +} + +void WeightLoader::fail(const char * what) { + std::fprintf(stderr, "vla(%s): %s\n", arch_, what); + ok_ = false; +} + +ggml_tensor * WeightLoader::fuse_typed(ggml_type want, const char * out_name, const std::vector & srcs) { + return fuse(want, out_name, srcs); +} + ggml_tensor * WeightLoader::fuse_gemm(const char * out_name, const std::vector & srcs) { return fuse(gemm_, out_name, srcs); } @@ -106,8 +172,17 @@ ggml_tensor * WeightLoader::fuse(ggml_type want, const char * out_name, const st return nullptr; } - const ggml_type rt = g_.resident_type(first, want); - const bool is1d = ggml_n_dims(first) == 1; + // A source registered with as_float() reads as an F32 [K, N] tensor. + struct View { ggml_type file, resident; int64_t ne0, ne1; int nd; }; + auto view = [&](const std::string & s, const ggml_tensor * gs) { + const auto it = as_float_.find(s); + if (it != as_float_.end()) + return View{GGML_TYPE_F32, float_type(want), it->second.K, it->second.N, 2}; + return View{gs->type, g_.resident_type(gs, want), gs->ne[0], gs->ne[1], ggml_n_dims(gs)}; + }; + const View v0 = view(srcs[0], first); + const ggml_type rt = v0.resident; + const bool is1d = v0.nd == 1; int64_t rows = 0; for (const std::string & s : srcs) { const ggml_tensor * gs = g_.meta(s.c_str()); @@ -116,16 +191,20 @@ ggml_tensor * WeightLoader::fuse(ggml_type want, const char * out_name, const st ok_ = false; return nullptr; } - if (g_.resident_type(gs, want) != rt || (!is1d && gs->ne[0] != first->ne[0])) { + // Same resident type and row shape; a tensor copied raw (a packed type, + // FoldQuant INT8 codes) also needs every source in that same type. + const View v = view(s, gs); + if (v.resident != rt || (rt == v0.file && v.file != v0.file) || + (!is1d && v.ne0 != v0.ne0) || v.nd != v0.nd) { std::fprintf(stderr, "vla(%s): %s does not match %s for fusing\n", arch_, s.c_str(), srcs[0].c_str()); ok_ = false; return nullptr; } - rows += is1d ? gs->ne[0] : gs->ne[1]; + rows += is1d ? v.ne0 : v.ne1; } ggml_tensor * t = is1d ? ggml_new_tensor_1d(ctx_, rt, rows) - : ggml_new_tensor_2d(ctx_, rt, first->ne[0], rows); + : ggml_new_tensor_2d(ctx_, rt, v0.ne0, rows); if (!t) { std::fprintf(stderr, "vla(%s): ggml_new_tensor failed for %s\n", arch_, out_name); ok_ = false; @@ -158,8 +237,8 @@ bool WeightLoader::upload(ggml_backend_t backend, ggml_backend_buffer_t * out_bu const bool gn = std::find(gemma_norms_.begin(), gemma_norms_.end(), name) != gemma_norms_.end(); - std::vector bytes = g_.read_convert(name, t->type, gn); - if (bytes.empty() || bytes.size() != ggml_nbytes(t)) { + std::vector bytes; + if (!read_resident(name, t->type, gn, bytes) || bytes.size() != ggml_nbytes(t)) { std::fprintf(stderr, "vla(%s): failed to load %s (got %zu bytes, expected %zu, type=%d)\n", arch_, name, bytes.size(), ggml_nbytes(t), (int) t->type); @@ -171,8 +250,8 @@ bool WeightLoader::upload(ggml_backend_t backend, ggml_backend_buffer_t * out_bu for (const Fused & f : fused_) { std::vector parts; for (const std::string & s : f.srcs) { - std::vector b = g_.read_convert(s.c_str(), f.dst->type); - if (b.empty()) { + std::vector b; + if (!read_resident(s, f.dst->type, false, b)) { std::fprintf(stderr, "vla(%s): fused fill: read %s failed\n", arch_, s.c_str()); return false; } diff --git a/src/loader.h b/src/loader.h index 649c9e9..87d09d0 100644 --- a/src/loader.h +++ b/src/loader.h @@ -24,6 +24,9 @@ #include "ggml-backend.h" #include +#include +#include +#include #include #include @@ -43,6 +46,7 @@ class WeightLoader { // Explicit resident type, for weights that are neither a plain GEMM input // nor F32 (BitVLA's int2-packed ternary blocks). ggml_tensor * typed(ggml_type want, const char * fmt, ...) __attribute__((format(printf, 3, 4))); + ggml_tensor * opt_typed(ggml_type want, const char * fmt, ...) __attribute__((format(printf, 3, 4))); // A miss is not an error. ggml_tensor * opt_gemm(const char * fmt, ...) __attribute__((format(printf, 2, 3))); @@ -55,6 +59,20 @@ class WeightLoader { // synthetic and need not exist in the file. ggml_tensor * fuse_gemm(const char * out_name, const std::vector & srcs); ggml_tensor * fuse_f32 (const char * out_name, const std::vector & srcs); + ggml_tensor * fuse_typed(ggml_type want, const char * out_name, const std::vector & srcs); + + // A tensor of another type in the file (a FoldQuant site's INT codes) that + // gemm/opt_gemm/fuse_gemm declare as a float [K, N] weight in the resident + // GEMM type; make() returns its N*K floats (row-major [N][K]) at upload. + using FloatMaker = std::function &)>; + void as_float(const std::string & name, int64_t K, int64_t N, FloatMaker make); + + // For helpers that inspect the file before declaring (foldquant.cpp). + gguf_reader & reader() { + return g_; + } + // Record a failure detected outside declare(); upload() then refuses. + void fail(const char * what); ggml_type gemm_type() const { return gemm_; @@ -68,6 +86,15 @@ class WeightLoader { private: ggml_tensor * declare(ggml_type want, bool required, bool gemma_norm, const char * fmt, va_list ap); ggml_tensor * fuse(ggml_type want, const char * out_name, const std::vector & srcs); + // Bytes of a resident tensor filled from GGUF tensor `name` (a registered + // float rebuild, or the file's tensor converted). + bool read_resident(const std::string & name, ggml_type type, bool gemma_norm, std::vector & out); + + struct AsFloat { + int64_t K = 0, N = 0; + FloatMaker make; + }; + std::map as_float_; const char * arch_; gguf_reader & g_; diff --git a/src/models/gr00tn1d5.cpp b/src/models/gr00tn1d5.cpp index c18b3bb..719f261 100644 --- a/src/models/gr00tn1d5.cpp +++ b/src/models/gr00tn1d5.cpp @@ -15,6 +15,7 @@ #include "arch.h" #include "options.h" #include "backend.h" +#include "foldquant.h" #include "gguf_reader.h" #include "layers/embed.h" #include "layers/linear.h" @@ -226,10 +227,13 @@ std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, (long long) m->action_horizon, (long long) m->action_dim, (long long) m->num_steps, (long long) m->aex.embodiment_id, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); + FoldQuantSpec fq = foldquant_parse(g, "gr00t_n1_5"); const Backend b = backend_init("vla(gr00tn1d5)", m->n_threads); if (!b.handle) return nullptr; m->backend = b.handle; + if (!foldquant_check_backend("vla(gr00tn1d5)", b, fq, opts.weight_dtype.has_value())) + return nullptr; ggml_init_params wp = { (size_t) 32*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); @@ -244,7 +248,7 @@ std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, m->mm_W = L.gemm("mm.fc.weight"); m->mm_b = L.f32 ("mm.fc.bias"); - m->lm.declare(L, "vlm"); + m->lm.declare(L, "vlm", fq.present ? &fq.llm : nullptr); m->vlln_w = L.f32("aex.vlln.weight"); m->vlln_b = L.f32("aex.vlln.bias"); @@ -253,7 +257,7 @@ std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, if (!m->aex.declare(L, g, m->backend, "aex")) return nullptr; m->future_tokens = L.f32("aex.future_tokens"); - m->dit.declare(L, "aex.dit"); + m->dit.declare(L, "aex.dit", false, false, nullptr, fq.present ? &fq.action : nullptr); if (!L.upload(m->backend, &m->weight_buf)) return nullptr; diff --git a/src/models/gr00tn1d6.cpp b/src/models/gr00tn1d6.cpp index 0824a39..5e9b131 100644 --- a/src/models/gr00tn1d6.cpp +++ b/src/models/gr00tn1d6.cpp @@ -15,6 +15,7 @@ #include "arch.h" #include "options.h" #include "backend.h" +#include "foldquant.h" #include "gguf_reader.h" #include "layers/embed.h" #include "layers/ffn.h" @@ -243,10 +244,13 @@ std::unique_ptr gr00t_n1_6_create(const std::string& mmproj_path, (long long) m->action_horizon, (long long) m->action_dim, (long long) m->max_state_dim, (long long) m->num_steps, (long long) m->aex.embodiment_id, m->matmul_type == GGML_TYPE_F32 ? "F32" : "BF16"); + FoldQuantSpec fq = foldquant_parse(g, "gr00t_n1_6"); const Backend b = backend_init("vla(gr00tn1d6)", m->n_threads); if (!b.handle) return nullptr; m->backend = b.handle; + if (!foldquant_check_backend("vla(gr00tn1d6)", b, fq, opts.weight_dtype.has_value())) + return nullptr; ggml_init_params wp = { (size_t) 32*1024*1024, nullptr, true }; m->ctx_weights = ggml_init(wp); @@ -266,14 +270,14 @@ std::unique_ptr gr00t_n1_6_create(const std::string& mmproj_path, m->mm_fc2_w = L.gemm("mm.fc2.weight"); m->mm_fc2_b = L.f32 ("mm.fc2.bias"); - m->lm.declare(L, "vlm"); + m->lm.declare(L, "vlm", fq.present ? &fq.llm : nullptr); m->vlln_w = L.f32("aex.vlln.weight"); m->vlln_b = L.f32("aex.vlln.bias"); if (!m->aex.declare(L, g, m->backend, "aex")) return nullptr; - m->dit.declare(L, "aex.dit"); + m->dit.declare(L, "aex.dit", false, false, nullptr, fq.present ? &fq.action : nullptr); if (!L.upload(m->backend, &m->weight_buf)) return nullptr; diff --git a/src/models/gr00tn1d7.cpp b/src/models/gr00tn1d7.cpp index 433f184..d0ce433 100644 --- a/src/models/gr00tn1d7.cpp +++ b/src/models/gr00tn1d7.cpp @@ -25,6 +25,7 @@ #include "ggml.h" #include "ggml-backend.h" #include "backend.h" +#include "env_flag.h" #include "gguf_reader.h" #include "scratch_ctx.h" #include "layers/embed.h" @@ -196,6 +197,7 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, } if (!load_config(g, opts, *m, m->cfg)) return nullptr; + FoldQuantSpec fq = foldquant_parse(g, "gr00t_n1_7"); std::printf("vla(gr00tn1d7): vit=Qwen3-VL %lldd×%lldL×%lldh (Conv3d patch %lld², temporal %lld; learned pos %lld + 2D rope; deepstack@{%lld,%lld,%lld}; merge÷%lld) " "lm=Qwen3-VL %lldd×%lldL (%lldq/%lldkv×%lld, θ=%g) vlsa=%lldL×%lldh×%lld dit=AlternateVLDiT %lldL×%lldh×%lld(inner %lld) attend_text_every_n=%lld " "in_emb=%lld horizon=%lld action_dim=%lld max_state=%lld N_steps=%lld embodiment=%lld resident=%s\n", @@ -213,6 +215,8 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, return nullptr; } m->backend = b.handle; + if (!foldquant_check_backend("vla(gr00tn1d7)", b, fq, opts.weight_dtype.has_value())) + return nullptr; } ggml_init_params wp = { (size_t) 32*1024*1024, nullptr, true }; @@ -225,7 +229,7 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, WeightLoader L("gr00tn1d7", g, m->ctx_weights, m->matmul_type); m->vit.declare(L, "vit"); - m->lm.declare(L, "vlm"); + m->lm.declare(L, "vlm", fq.present ? &fq.llm : nullptr); m->vlln_w = L.f32("aex.vlln.weight"); m->vlln_b = L.f32("aex.vlln.bias"); @@ -233,7 +237,18 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, if (!m->aex.declare(L, g, m->backend, "aex")) return nullptr; - m->dit.declare(L, "aex.dit", true, m->dit_interleave != 0); + m->dit.declare(L, "aex.dit", true, m->dit_interleave != 0, nullptr, fq.present ? &fq.action : nullptr); + if (fq.present) { + // Execution order of the FoldQuant GEMMs, so each one can prefetch the + // next site's weights (cross-attention K/V run before the step loop and + // are left out of the chain). + std::vector order; + for (auto & b : m->lm.blk) + for (FqLinear * s : {&b.fq_q, &b.fq_k, &b.fq_v, &b.fq_o, &b.fq_gate, &b.fq_up, &b.fq_down}) order.push_back(s); + for (auto & b : m->dit.blk) + for (FqLinear * s : {&b.fq_qkv, &b.fq_q, &b.fq_o, &b.fq_ff0, &b.fq_ff2}) order.push_back(s); + fq_link_prefetch(order); + } if (!L.upload(m->backend, &m->weight_buf)) return nullptr; @@ -396,6 +411,38 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { ggml_backend_tensor_set(gio.t_txt_idx, prompt.text_pos.data(), 0, ggml_nbytes(gio.t_txt_idx)); graph_unique_names(gf); + // VLA_GRAPH_DEBUG=1: report what changes between two computes of the cached + // graph, mirroring ggml-cuda's graph-reuse test (whole node struct + source + // data pointers / shapes); any change there defeats CUDA-graph replay. + if (env_flag("VLA_GRAPH_DEBUG")) { + struct Prop { ggml_tensor t; const void * sp[GGML_MAX_SRC]; }; + static std::vector prev; + const int n = ggml_graph_n_nodes(gf); + std::vector cur((size_t) n); + for (int i = 0; i < n; ++i) { + ggml_tensor * t = ggml_graph_node(gf, i); + std::memcpy(&cur[(size_t) i].t, t, sizeof(ggml_tensor)); + for (int j = 0; j < GGML_MAX_SRC; ++j) cur[(size_t) i].sp[j] = t->src[j] ? t->src[j]->data : nullptr; + } + if (prev.size() == cur.size()) { + int changed = 0; + for (int i = 0; i < n; ++i) { + const ggml_tensor & a = prev[(size_t) i].t, & b = cur[(size_t) i].t; + if (std::memcmp(&a, &b, sizeof(ggml_tensor)) != 0 || std::memcmp(prev[(size_t) i].sp, cur[(size_t) i].sp, sizeof(cur[(size_t) i].sp)) != 0) { + if (changed < 5) + std::printf("vla(graph-debug): node %d %s changed: data %p->%p op_params %d name %d flags %d extra %p->%p srcdata %d\n", + i, ggml_get_name(&b), a.data, b.data, std::memcmp(a.op_params, b.op_params, sizeof(a.op_params)) != 0, + std::strcmp(a.name, b.name) != 0, a.flags != b.flags, a.extra, b.extra, + std::memcmp(prev[(size_t) i].sp, cur[(size_t) i].sp, sizeof(cur[(size_t) i].sp)) != 0); + ++changed; + } + } + std::printf("vla(graph-debug): %d of %d nodes changed since the previous compute\n", changed, n); + } else if (!prev.empty()) { + std::printf("vla(graph-debug): graph size changed %zu -> %zu\n", prev.size(), cur.size()); + } + prev = std::move(cur); + } const auto tc0 = std::chrono::steady_clock::now(); const ggml_status st = ggml_backend_graph_compute(backend, gf); const auto tc1 = std::chrono::steady_clock::now(); diff --git a/src/models/pi05.cpp b/src/models/pi05.cpp index 7a63218..5992b0c 100644 --- a/src/models/pi05.cpp +++ b/src/models/pi05.cpp @@ -26,6 +26,8 @@ #include "layers/linear.h" #include "modules/preprocess.h" #include "modules/prompt.h" +#include "foldquant.h" +#include "layers/fq_linear.h" #include #include @@ -107,14 +109,16 @@ ggml_tensor * build_expert_layer( ggml_tensor * gate_attn = nullptr; ggml_tensor * x_norm = build_adarms(ctx, x_in, mod_attn, h, cfg.rms_eps, &gate_attn); + // FoldQuant: the act nodes quantize the adaRMS outputs; the residuals are + // gated (out * gate + x), so they stay out of the GEMM epilogues. ggml_tensor * o_out = gemma_attn(ctx, w, x_norm, positions, cfg, seq, cached_K, cached_V, nullptr, - nullptr, nullptr, GGML_TYPE_F32, false); + nullptr, nullptr, GGML_TYPE_F32, false, x_norm); ggml_tensor * h1 = ggml_add(ctx, x_in, ggml_mul(ctx, o_out, gate_attn)); ggml_tensor * gate_ffn = nullptr; ggml_tensor * x_norm_mlp = build_adarms(ctx, h1, mod_ffn, h, cfg.rms_eps, &gate_ffn); - ggml_tensor * mlp_out = gemma_mlp(ctx, w, x_norm_mlp, GGML_TYPE_F32); + ggml_tensor * mlp_out = gemma_mlp(ctx, w, x_norm_mlp, GGML_TYPE_F32, x_norm_mlp); return ggml_add(ctx, h1, ggml_mul(ctx, mlp_out, gate_ffn)); } @@ -247,6 +251,7 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, if (!load_pi_config(g, ckpt_path, 0, m->cfg) || !resolve_num_steps("pi05", opts, m->cfg.num_steps)) return nullptr; const Config & cfg = m->cfg; + FoldQuantSpec fq = foldquant_parse(g, "pi05"); m->quantile_norm = g.has("pi05.norm_mode") && g.str("pi05.norm_mode") == "quantiles"; std::printf("vla(pi05): hidden=%lld inter=%lld heads=%lldq/%lldkv x%lld n_layers=%lld " "expert_h=%lld expert_inter=%lld chunk=%lld steps=%d real_state=%lld real_action=%lld " @@ -263,6 +268,8 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, return nullptr; } m->backend = b.handle; + if (!foldquant_check_backend("vla(pi05)", b, fq, opts.weight_dtype.has_value())) + return nullptr; } // The SigLIP tower is now bundled in the ckpt GGUF; mmproj_path is ignored. @@ -282,18 +289,39 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, m->vis.declare(L); - m->pl.declare(L, "vlm", cfg.n_layers, false); + m->pl.declare(L, "vlm", cfg.n_layers, false, fq.present ? &fq.llm : nullptr, cfg.rms_eps); m->ex_layers.resize(cfg.n_layers); for (int64_t i=0; iex_layers[i]; - w.Wq = L.gemm("aex.blk.%lld.attn_q.weight", (long long)i); - w.Wk = L.gemm("aex.blk.%lld.attn_k.weight", (long long)i); - w.Wv = L.gemm("aex.blk.%lld.attn_v.weight", (long long)i); - w.Wo = L.gemm("aex.blk.%lld.attn_o.weight", (long long)i); - w.Wgate = L.gemm("aex.blk.%lld.ffn_gate.weight", (long long)i); - w.Wup = L.gemm("aex.blk.%lld.ffn_up.weight", (long long)i); - w.Wdown = L.gemm("aex.blk.%lld.ffn_down.weight", (long long)i); + const long long ii = (long long) i; + if (fq.present) { + // FoldQuant sites (action recipe): the adaRMS output is quantized per + // site group with the shipped SmoothQuant vector (ascale), no gamma. + const FqModuleSpec & a = fq.action; + w.fq_q = fq_declare_linear(L, a, "qkv", false, nullptr, 0.0f, "aex.blk.%lld.attn_q", ii); + w.fq_k = fq_declare_linear(L, a, "qkv", false, nullptr, 0.0f, "aex.blk.%lld.attn_k", ii); + w.fq_v = fq_declare_linear(L, a, "qkv", false, nullptr, 0.0f, "aex.blk.%lld.attn_v", ii); + w.fq_o = fq_declare_linear(L, a, "o", false, nullptr, 0.0f, "aex.blk.%lld.attn_o", ii); + w.fq_gate = fq_declare_linear(L, a, "gateup", false, nullptr, 0.0f, "aex.blk.%lld.ffn_gate", ii); + w.fq_up = fq_declare_linear(L, a, "gateup", false, nullptr, 0.0f, "aex.blk.%lld.ffn_up", ii); + w.fq_down = fq_declare_linear(L, a, "down", false, nullptr, 0.0f, "aex.blk.%lld.ffn_down", ii); + if (!w.fq_q != !w.fq_k || !w.fq_q != !w.fq_v || !w.fq_gate != !w.fq_up) + L.fail("FoldQuant: an expert layer's q/k/v (and gate/up) must all be INT or all float"); + } + if (!w.fq_q) { + w.Wq = L.gemm("aex.blk.%lld.attn_q.weight", ii); + w.Wk = L.gemm("aex.blk.%lld.attn_k.weight", ii); + w.Wv = L.gemm("aex.blk.%lld.attn_v.weight", ii); + } + if (!w.fq_o) + w.Wo = L.gemm("aex.blk.%lld.attn_o.weight", ii); + if (!w.fq_gate) { + w.Wgate = L.gemm("aex.blk.%lld.ffn_gate.weight", ii); + w.Wup = L.gemm("aex.blk.%lld.ffn_up.weight", ii); + } + if (!w.fq_down) + w.Wdown = L.gemm("aex.blk.%lld.ffn_down.weight", ii); } m->W_ain = L.f32("action_in_proj.weight"); m->b_ain = L.f32("action_in_proj.bias"); diff --git a/src/modules/dit_head.cpp b/src/modules/dit_head.cpp index addee52..b8051a3 100644 --- a/src/modules/dit_head.cpp +++ b/src/modules/dit_head.cpp @@ -18,6 +18,7 @@ #include "layers/attn.h" #include "layers/embed.h" #include "layers/ffn.h" +#include "layers/fq_linear.h" #include "layers/linear.h" #include "ggml-alloc.h" @@ -28,7 +29,8 @@ namespace vla { -void DitHead::declare(WeightLoader & L, const char * prefix, bool fuse_qkv, bool interleave, const char * outer) { +void DitHead::declare(WeightLoader & L, const char * prefix, bool fuse_qkv, bool interleave, const char * outer, + const FqModuleSpec * fq) { if (!outer) outer = prefix; @@ -40,22 +42,52 @@ void DitHead::declare(WeightLoader & L, const char * prefix, bool fuse_qkv, bool blk.resize(cfg.layers); for (int64_t i=0; ine[1]; - if (w.Wkv) { - ggml_tensor * kvp = linear(C, w.Wkv, w.bkv, src); + if (w.Wkv || w.fq_kv) { + ggml_tensor * kvp = w.fq_kv ? fq_linear(C, w.fq_kv, src) : linear(C, w.Wkv, w.bkv, src); + if (w.fq_kv && w.fq_kv.gemm.heads) { // written in head layout by the epilogue + *K_out = fq_head_view(C, kvp, w.fq_kv, 0, Tkv); + *V_out = fq_head_view(C, kvp, w.fq_kv, 1, Tkv); + return; + } *K_out = ggml_cont(C, ggml_permute(C, head_view(C, kvp, hd, heads, Tkv, cfg.hidden, 2, 0), 0, 2, 1, 3)); *V_out = ggml_cont(C, ggml_permute(C, head_view(C, kvp, hd, heads, Tkv, cfg.hidden, 2, 1), 1, 2, 0, 3)); return; } + if (w.fq_k) { + ggml_tensor * xq = xq_pre ? xq_pre : fq_act(C, w.fq_k, src); + ggml_tensor * kp = fq_gemm(C, w.fq_k, xq), * vp = fq_gemm(C, w.fq_v, xq); + *K_out = w.fq_k.gemm.heads ? fq_head_view(C, kp, w.fq_k, 0, Tkv) : to_heads (C, kp, hd, heads, Tkv); + *V_out = w.fq_v.gemm.heads ? fq_head_view(C, vp, w.fq_v, 0, Tkv) : to_heads_v(C, vp, hd, heads, Tkv); + return; + } *K_out = to_heads (C, linear(C, w.Wk, w.bk, src), hd, heads, Tkv); *V_out = to_heads_v(C, linear(C, w.Wv, w.bv, src), hd, heads, Tkv); } @@ -117,25 +184,42 @@ ggml_tensor * DitHead::block(ggml_context * C, const DitLayerW & w, ggml_tensor ggml_tensor * xn = ggml_norm(C, h, cfg.ln_eps); ggml_tensor * n = ggml_add(C, ggml_add(C, xn, ggml_mul(C, xn, sc)), sh); ggml_tensor *Q, *K, *V; - if (!enc && w.Wqkv) { - ggml_tensor * qkv = linear(C, w.Wqkv, w.bqkv, n); + if (!enc && (w.Wqkv || w.fq_qkv)) { + ggml_tensor * qkv = w.fq_qkv ? fq_linear(C, w.fq_qkv, n) : linear(C, w.Wqkv, w.bqkv, n); + if (w.fq_qkv && w.fq_qkv.gemm.heads) { + Q = fq_head_view(C, qkv, w.fq_qkv, 0, Tk); + K = fq_head_view(C, qkv, w.fq_qkv, 1, Tk); + V = fq_head_view(C, qkv, w.fq_qkv, 2, Tk); + } else { Q = ggml_cont(C, ggml_permute(C, head_view(C, qkv, hd, heads, Tk, dim, 3, 0), 0, 2, 1, 3)); K = ggml_cont(C, ggml_permute(C, head_view(C, qkv, hd, heads, Tk, dim, 3, 1), 0, 2, 1, 3)); V = ggml_cont(C, ggml_permute(C, head_view(C, qkv, hd, heads, Tk, dim, 3, 2), 1, 2, 0, 3)); + } } else { - Q = to_heads(C, linear(C, w.Wq, w.bq, n), hd, heads, Tk); + // Self-attention with separate FoldQuant q/k/v: one activation blob of n + // serves all three projections. + ggml_tensor * xq = (w.fq_q && !enc && !K_pre && w.fq_k) ? fq_act(C, w.fq_q, n) : nullptr; + if (w.fq_q) { + ggml_tensor * qp = xq ? fq_gemm(C, w.fq_q, xq) : fq_linear(C, w.fq_q, n); + Q = w.fq_q.gemm.heads ? fq_head_view(C, qp, w.fq_q, 0, Tk) : to_heads(C, qp, hd, heads, Tk); + } else { + Q = to_heads(C, linear(C, w.Wq, w.bq, n), hd, heads, Tk); + } if (K_pre) { K = K_pre; V = V_pre; } else { - kv(C, w, enc ? enc : n, &K, &V); + kv(C, w, enc ? enc : n, &K, &V, xq); } } ggml_tensor * att = attention(C, Q, K, V, nullptr, scale, dim, Tk); - ggml_tensor * h1 = ggml_add(C, h, linear(C, w.Wo, w.bo, att)); + // FoldQuant: both residual adds ride in the GEMM epilogues. + ggml_tensor * h1 = w.fq_o ? fq_linear(C, w.fq_o, att, h) : ggml_add(C, h, linear(C, w.Wo, w.bo, att)); ggml_tensor * n3 = ggml_norm(C, h1, cfg.ln_eps); + if (w.fq_ff0) + return fq_linear(C, w.fq_ff2, ggml_gelu(C, fq_linear(C, w.fq_ff0, n3)), h1); return ggml_add(C, h1, ffn_gelu(C, w.Wff0, w.bff0, w.Wff2, w.bff2, n3)); } diff --git a/src/modules/dit_head.h b/src/modules/dit_head.h index 9a003e2..fc63d44 100644 --- a/src/modules/dit_head.h +++ b/src/modules/dit_head.h @@ -18,6 +18,7 @@ #pragma once +#include "foldquant.h" #include "loader.h" #include "ggml.h" @@ -34,6 +35,10 @@ struct DitLayerW { ggml_tensor *Wff0, *bff0, *Wff2, *bff2; ggml_tensor *Wqkv = nullptr, *bqkv = nullptr, *Wkv = nullptr, *bkv = nullptr; + + // FoldQuant sites (docs/QUANTIZATION.md); empty where the file ships a float + // GEMM. fq_qkv / fq_kv are the fused forms of fq_q/fq_k/fq_v. + FqLinear fq_q, fq_k, fq_v, fq_o, fq_ff0, fq_ff2, fq_qkv, fq_kv; }; struct DitCfg { @@ -53,11 +58,13 @@ struct DitHead { // outer names time_emb and proj_out when they do not sit under the block // prefix; null means they do. + // fq: the action module's FoldQuant parameters when the GGUF carries them. void declare(WeightLoader & L, const char * prefix, bool fuse_qkv = false, bool interleave = false, - const char * outer = nullptr); + const char * outer = nullptr, const FqModuleSpec * fq = nullptr); + // xq_pre: an activation blob already computed from src (shared with q). void kv(ggml_context * C, const DitLayerW & w, ggml_tensor * src, - ggml_tensor ** K_out, ggml_tensor ** V_out) const; + ggml_tensor ** K_out, ggml_tensor ** V_out, ggml_tensor * xq_pre = nullptr) const; ggml_tensor * block(ggml_context * C, const DitLayerW & w, ggml_tensor * h, ggml_tensor * mod, ggml_tensor * enc, ggml_tensor * K_pre = nullptr, ggml_tensor * V_pre = nullptr) const; diff --git a/src/modules/gemma_expert.h b/src/modules/gemma_expert.h index 65b3041..f14001b 100644 --- a/src/modules/gemma_expert.h +++ b/src/modules/gemma_expert.h @@ -22,6 +22,8 @@ #include "backend.h" #include "gguf_reader.h" #include "layers/norm.h" +#include "foldquant.h" +#include "layers/fq_linear.h" #include "loader.h" #include "model.h" #include "modules/siglip_vit.h" @@ -46,25 +48,48 @@ struct GemmaLayerW { ggml_tensor * Wgate = nullptr; ggml_tensor * Wup = nullptr; ggml_tensor * Wdown = nullptr; + // FoldQuant sites (pi0.5 prefix tower). The RMSNorm and its folded gamma + // (ln_in / ln_post, loaded as 1 + w) ride in the q/k/v and gate/up act nodes. + FqLinear fq_q, fq_k, fq_v, fq_o, fq_gate, fq_up, fq_down; }; struct GemmaStack { std::vector blk; ggml_tensor * output_norm = nullptr; - void declare(WeightLoader & L, const char * prefix, int64_t layers, bool with_output_norm) { + // fq: the LLM FoldQuant spec when the GGUF carries one (pi0.5), else null. + void declare(WeightLoader & L, const char * prefix, int64_t layers, bool with_output_norm, + const FqModuleSpec * fq = nullptr, float rms_eps = 1e-6f) { blk.resize(layers); for (int64_t i=0; i namespace vla { -void Qwen3LM::declare(WeightLoader & L, const char * prefix) { +void Qwen3LM::declare(WeightLoader & L, const char * prefix, const FqModuleSpec * fq) { output_norm = L.f32("%s.output_norm.weight", prefix); blk.resize(cfg.layers); for (int64_t i=0; i blk; ggml_tensor * output_norm = nullptr; - void declare(WeightLoader & L, const char * prefix); + // fq: the LLM module's FoldQuant parameters when the GGUF carries them, + // null otherwise (every site is then a float GEMM and the graph is unchanged). + void declare(WeightLoader & L, const char * prefix, const FqModuleSpec * fq = nullptr); ggml_tensor * block(ggml_context * C, const Qwen3LayerW & w, ggml_tensor * h, ggml_tensor * pos, ggml_tensor * mask, int64_t seq) const; diff --git a/src/openvino/foldquant_ov.cpp b/src/openvino/foldquant_ov.cpp new file mode 100644 index 0000000..14a4c4d --- /dev/null +++ b/src/openvino/foldquant_ov.cpp @@ -0,0 +1,232 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant's two GGML_OP_CUSTOM nodes (src/layers/fq_linear.h) translated into +// OpenVINO ops, so a FoldQuant GGUF runs natively on ggml's OpenVINO backend. +// Compiled into the ggml-openvino target by vla.cpp's CMakeLists; the two hook +// points (the op table entry and supports_op) are added by +// scripts/patch_ggml_openvino.py. +// +// fq_act (x[, gamma][, ascale]) -> [.., T, K] F32: the integer-valued codes times +// the per-token scale. ggml declares an I8 blob for this node; only fq_gemm +// reads it, so the translator is free to carry it as floats. +// fq_gemm (w, act[, bias][, residual]) -> [.., T, N] F32. +// +// The arithmetic is the CPU reference's (src/foldquant_ref.cpp): RMSNorm with the +// folded gamma, the ascale divide before or after the block-normalised +// Sylvester-Hadamard rotation, scale = max(clip * amax / qmax, 1e-12), codes = +// clamp(round_half_even(y * (1 / scale)), -qmax, qmax). The weights stay INT8 or +// INT4 constants (W4 reinterprets the nibble bytes in place: OpenVINO's i4 packs +// the even element in the low nibble, as FoldQuant does) and are dequantized by +// wscale in the decompression pattern the plugins keep compressed. + +#include "openvino/node_context.h" +#include "openvino/op_table.h" +#include "openvino/utils.h" + +#include "foldquant.h" +#include "ggml-impl.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +const void * fq_userdata(const ggml_tensor * op) { + if (op == nullptr || op->op != GGML_OP_CUSTOM) + return nullptr; + const auto * p = reinterpret_cast(op->op_params); + return p->userdata; +} + +uint32_t fq_magic(const void * userdata) { + uint32_t m = 0; + if (userdata) + std::memcpy(&m, userdata, sizeof(m)); + return m; +} + +} // namespace + +// supports_op hook (declared again in the patched ggml-openvino.cpp): FoldQuant's +// nodes, and no other GGML_OP_CUSTOM. +bool vla_foldquant_ov_supports(const ggml_tensor * op); +bool vla_foldquant_ov_supports(const ggml_tensor * op) { + const uint32_t m = fq_magic(fq_userdata(op)); + if (m == vla::FQ_ACT_MAGIC) { + const auto * s = static_cast(fq_userdata(op)); + return op->src[0] && op->src[0]->type == GGML_TYPE_F32 && s->rot_block <= 64; + } + if (m == vla::FQ_GEMM_MAGIC) { + const auto * s = static_cast(fq_userdata(op)); + // The head-laid-out epilogue writes a layout the graph reads through views + // of raw memory; the vla.cpp side turns it off on this backend. + return s->heads == 0 && op->src[0] && op->src[0]->type == GGML_TYPE_I8; + } + return false; +} + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +namespace { + +using ov::op::v0::Constant; + +Output f32_scalar(float v) { + return Constant::create(element::f32, Shape{1}, {v}); +} + +Output i64_vec(std::vector v) { + return Constant::create(element::i64, Shape{v.size()}, v); +} + +// Normalised natural-order Sylvester-Hadamard matrix of size n (symmetric). +Output hadamard(int n) { + std::vector h((size_t) n * n); + const float norm = 1.0f / std::sqrt((float) n); + for (int i = 0; i < n; ++i) + for (int j = 0; j < n; ++j) + h[(size_t) i * n + j] = (__builtin_popcount((unsigned) (i & j)) & 1) ? -norm : norm; + return Constant::create(element::f32, Shape{(size_t) n, (size_t) n}, h); +} + +// A [K] / [1, K] weight vector shaped to broadcast along the last axis. +Output row_vector(const Output & v, int64_t K) { + Output r = v; + if (r.get_element_type() != element::f32) + r = std::make_shared(r, element::f32); + return std::make_shared(r, i64_vec({K}), false); +} + +OutputVector translate_act(const NodeContext & context, const vla::FqActSpec & s) { + const int64_t K = s.K; + Output x = process_view_input_new(context, 0); + if (x.get_element_type() != element::f32) + x = std::make_shared(x, element::f32); + size_t next = 1; + Output y = x; + if (s.has_gamma) { + const Output gamma = row_vector(context.get_input((int) next++), K); + auto ms = std::make_shared(std::make_shared(x, x), + i64_vec({-1}), true); + auto rstd = std::make_shared( + f32_scalar(1.0f), std::make_shared(std::make_shared(ms, f32_scalar(s.eps)))); + y = std::make_shared(std::make_shared(x, rstd), gamma); + } + Output a; + if (s.has_ascale) + a = row_vector(context.get_input((int) next++), K); + if (s.has_ascale && s.fold_before) + y = std::make_shared(y, a); + if (s.rot_block > 1) { + // Rotate each rot_block chunk of the last axis: [.., K] -> [.., K/rb, rb] . H -> [.., K]. + auto shape = std::make_shared(y, element::i64); + auto blk = std::make_shared(y, i64_vec({-1, K / s.rot_block, s.rot_block}), false); + auto rot = std::make_shared(blk, hadamard(s.rot_block), false, false); + y = std::make_shared(rot, shape, false); + } + if (s.has_ascale && !s.fold_before) + y = std::make_shared(y, a); + + const float qmax = s.abits == 4 ? 7.0f : 127.0f; + auto amax = std::make_shared(std::make_shared(y), i64_vec({-1}), true); + auto scale = std::make_shared( + std::make_shared(std::make_shared(amax, f32_scalar(s.clip)), + f32_scalar(qmax)), + f32_scalar(1e-12f)); + auto inv = std::make_shared(f32_scalar(1.0f), scale); + auto q = std::make_shared( + std::make_shared(std::make_shared(y, inv), + ov::op::v5::Round::RoundMode::HALF_TO_EVEN), + -qmax, qmax); + // codes * scale: the per-token scale factors out of the GEMM, so carrying it + // inside the activation saves the codes/scale split at every consumer. + auto out = std::make_shared(q, scale); + return rename_outputs_with_suffix({out}, context.get_name()); +} + +OutputVector translate_gemm(const NodeContext & context, const vla::FqGemmSpec & s) { + const int64_t K = s.K, N = s.N; + + // INT codes [N, K_pack] as the backend's weight constant; W4 reinterprets the + // same bytes as an i4 [N, K] constant (the weight buffer outlives the model). + auto wnode = std::dynamic_pointer_cast(context.get_input(0).get_node_shared_ptr()); + OPENVINO_ASSERT(wnode, "FoldQuant: weight of ", context.get_name(), " is not a constant"); + std::shared_ptr w = wnode; + if (s.wbits == 4) { + ov::Tensor t(element::i4, Shape{(size_t) N, (size_t) K}, const_cast(wnode->get_data_ptr())); + w = std::make_shared(t); + } else if (wnode->get_shape() != Shape{(size_t) N, (size_t) K}) { + w = std::make_shared(wnode, i64_vec({N, K}), false); + } + // Decompression: Convert(INT) * wscale per output row. + auto wscale = std::make_shared(row_vector(context.get_input(2), N), i64_vec({N, 1}), false); + auto wf = std::make_shared(std::make_shared(w, element::f32), wscale); + + // The activation arrives as codes * per-token scale (see translate_act). + Output y = std::make_shared(context.get_input(1), wf, false, true); + + // Optional sources after (w, act, wscale): the bias [N] and/or the residual + // shaped like the output; ggml drops null sources, so tell them by size. + for (size_t i = 3; i < context.get_input_size(); ++i) { + const Output extra = context.get_input((int) i); + const PartialShape ps = context.get_input_shape(i); + const bool is_bias = ps.is_static() && ov::shape_size(ps.to_shape()) == (size_t) N; + Output e = is_bias ? row_vector(extra, N) : extra; + if (e.get_element_type() != element::f32) + e = std::make_shared(e, element::f32); + y = std::make_shared(y, e); + } + return rename_outputs_with_suffix({y}, context.get_name()); +} + +} // namespace + +OutputVector translate_vla_foldquant(const NodeContext & context) { + const int32_t * params = context.get_output_op_params(); + const void * ud = reinterpret_cast(params)->userdata; + switch (fq_magic(ud)) { + case vla::FQ_ACT_MAGIC: + return translate_act(context, *static_cast(ud)); + case vla::FQ_GEMM_MAGIC: + return translate_gemm(context, *static_cast(ud)); + default: + OPENVINO_THROW("GGML_OP_CUSTOM ", context.get_name(), " is not a FoldQuant node"); + } +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/src/sycl/vla_sycl_foldquant.cpp b/src/sycl/vla_sycl_foldquant.cpp new file mode 100644 index 0000000..d25dbb0 --- /dev/null +++ b/src/sycl/vla_sycl_foldquant.cpp @@ -0,0 +1,845 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant's two custom nodes (src/layers/fq_linear.h) on ggml's SYCL backend, +// reached through the hook scripts/patch_ggml_sycl_ext_hook.py adds; the SYCL +// counterpart of src/cuda/vla_cuda_foldquant.cu and src/kernels/foldquant/. +// +// fq_act: the reduction tree foldquant_ref.h mirrors, with lane l of a 32-wide +// sub-group owning the 64-element chunks l, l+32, ... of a row, so the codes and +// scales are bit-identical to the CPU reference. Rows up to K = 8192 (all a VLA +// has but its 16384-wide down projections) spread over a work-group; longer ones +// run one sub-group per row. +// +// fq_gemm: INT8/INT4 x INT8/INT4 -> INT32, then ((float) acc * act_scale[m]) * +// wscale[n] (+ bias[n]) (+ residual), the reference's epilogue order. The sums +// are integer, so any tiling gives the reference's result: oneDNN's int8 matmul +// when ggml-sycl is built with it (GGML_SYCL_DNN, the default), else this file's +// GEMV (M <= 32) and XMX kernels. +// +// Exactness rests on the build flags in CMakeLists.txt: -ffp-contract=off for +// the source, and -cl-fp32-correctly-rounded-divide-sqrt for the device JIT, +// without which the driver's division and sqrt are approximate and move codes +// across rounding ties. test_foldquant_sycl_op checks every path byte for byte. + +#include "sycl/vla_sycl_foldquant.h" + +#include "foldquant.h" +#include "foldquant_ref.h" +#include "env_flag.h" +#include "ggml.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef VLA_FQ_SYCL_DNNL +#include +#include +#endif + +extern "C" { +typedef bool (*ggml_sycl_ext_forward_t)(struct ggml_tensor * dst, void * queue); +typedef bool (*ggml_sycl_ext_supports_t)(const struct ggml_tensor * op); +extern ggml_sycl_ext_forward_t ggml_sycl_ext_forward; +extern ggml_sycl_ext_supports_t ggml_sycl_ext_supports; +} + +namespace vla { + +namespace { + +constexpr int CHUNK = 64; +constexpr int SG = 32; // sub-group width: the reference's lane count +constexpr int ROWS = 4; // rows (sub-groups) per work-group + +struct ActParams { + const float * x; // [M][x_stride] + int64_t x_stride; + const float * ascale; // [K] or null + const float * gamma; // [K] or null + int8_t * blob; // [M][row_bytes] + int64_t row_bytes; + int64_t M, K; + bool fold_before; + float clip, eps, inv_sqrt_bs; +}; + +struct GemmParams { + const int8_t * w; // [N][K_pack] + const int8_t * blob; // [M][row_bytes] + const float * wscale; // [N] + const float * bias; // [N] or null + const float * res; // [M][N] or null + float * y; // [M][N] + int64_t M, N, K, row_bytes; +}; + +const void * fq_userdata(const ggml_tensor * op) { + if (op == nullptr || op->op != GGML_OP_CUSTOM) return nullptr; + struct custom_params { void * fun; int n_tasks; void * userdata; }; // ggml_custom_op_params + custom_params p; + std::memcpy(&p, op->op_params, sizeof(p)); + return p.userdata; +} + +uint32_t fq_magic(const void * ud) { + uint32_t m = 0; + if (ud) std::memcpy(&m, ud, sizeof(m)); + return m; +} + +// Load one 64-float chunk, apply the pre-quant chain and rotate it. ROT is a +// compile-time constant so every index into v[] is static (registers). +template +inline void load_chunk(const ActParams & a, const float * x, int64_t c0, float rstd, float * v) { + for (int i = 0; i < CHUNK; ++i) v[i] = x[c0 + i]; + if (a.gamma) + for (int i = 0; i < CHUNK; ++i) v[i] = (v[i] * rstd) * a.gamma[c0 + i]; + if (a.ascale && a.fold_before) + for (int i = 0; i < CHUNK; ++i) v[i] = v[i] / a.ascale[c0 + i]; + if (ROT > 1) { + // Natural-order Sylvester butterfly on every ROT sub-block of the chunk, + // the pairs and stage order of fwht_row. + for (int h = 1; h < ROT; h <<= 1) + for (int p = 0; p < CHUNK / 2; ++p) { + const int j = ((p / h) * 2 * h) + (p % h); + const float u = v[j], w = v[j + h]; + v[j] = u + w; + v[j + h] = u - w; + } + for (int i = 0; i < CHUNK; ++i) v[i] = v[i] * a.inv_sqrt_bs; + } + if (a.ascale && !a.fold_before) + for (int i = 0; i < CHUNK; ++i) v[i] = v[i] / a.ascale[c0 + i]; +} + +template +void act_kernel(const ActParams a, const sycl::nd_item<1> & it) { + const auto sg = it.get_sub_group(); + const int lane = (int) sg.get_local_linear_id(); + const int64_t m = (int64_t) it.get_group(0) * ROWS + (int64_t) sg.get_group_linear_id(); + if (m >= a.M) return; + const int64_t K = a.K; + const int nchunks = (int) (K / CHUNK); + const float * x = a.x + m * a.x_stride; + + float rstd = 0.0f; + if (a.gamma) { + float p = 0.0f; + for (int c = lane; c < nchunks; c += SG) + for (int i = 0; i < CHUNK; ++i) { + const float q = x[(int64_t) c * CHUNK + i]; + p = p + (q * q); + } + for (int off = 16; off > 0; off >>= 1) p = p + sycl::permute_group_by_xor(sg, p, off); + rstd = 1.0f / sycl::sqrt((p / (float) K) + a.eps); + } + + float v[CHUNK]; + float mx = 0.0f; + for (int c = lane; c < nchunks; c += SG) { + load_chunk(a, x, (int64_t) c * CHUNK, rstd, v); + for (int i = 0; i < CHUNK; ++i) mx = sycl::fmax(mx, sycl::fabs(v[i])); + } + for (int off = 16; off > 0; off >>= 1) mx = sycl::fmax(mx, sycl::permute_group_by_xor(sg, mx, off)); + + const float qmax = ABITS == 4 ? 7.0f : 127.0f; + float scale = (a.clip * mx) / qmax; + if (scale < 1e-12f) scale = 1e-12f; + // Multiply by the reciprocal, as the reference and the TensorRT kernels do. + const float inv = 1.0f / scale; + + int8_t * row = a.blob + m * a.row_bytes; + for (int c = lane; c < nchunks; c += SG) { + if (nchunks > SG) load_chunk(a, x, (int64_t) c * CHUNK, rstd, v); + if (ABITS == 8) { + int8_t * dst = row + (int64_t) c * CHUNK; + for (int i = 0; i < CHUNK; ++i) { + float q = sycl::rint(v[i] * inv); + q = sycl::fmin(qmax, sycl::fmax(-qmax, q)); + dst[i] = (int8_t) (int) q; + } + } else { + int8_t * dst = row + (int64_t) c * (CHUNK / 2); + for (int i = 0; i < CHUNK; i += 2) { + float q0 = sycl::rint(v[i] * inv), q1 = sycl::rint(v[i + 1] * inv); + q0 = sycl::fmin(qmax, sycl::fmax(-qmax, q0)); + q1 = sycl::fmin(qmax, sycl::fmax(-qmax, q1)); + dst[i / 2] = (int8_t) (((uint32_t) (int) q0 & 0xFu) | (((uint32_t) (int) q1 & 0xFu) << 4)); + } + } + } + if (lane == 0) { + float * s = (float *) (row + (ABITS == 8 ? K : K / 2)); + *s = scale; + } +} + +// Work-group-per-row variant for the few rows a VLA quantizes (10-41 action +// tokens, a few hundred LLM tokens): one sub-group per row leaves the GPU nearly +// idle. Here RW sub-groups share a row; sub-group w takes chunks w, w + RW, ... +// and lane l owns the element pair (2l, 2l + 1) of each chunk. The 64-wide +// butterfly runs stage h = 1 inside the lane and the other stages across lanes, +// forming exactly the (a + b, a - b) pairs fwht_row forms, and the RMSNorm sum +// of squares keeps the reference's lane-partial order (sub-group 0 runs the +// per-row loop), so the codes are still bit-identical. The CUDA path's +// act_row_kernel, ported. +constexpr int RW = 8, RT = RW * SG, MAXC = 16; // K <= RW * MAXC * 64 + +template +inline void rot_pair(const sycl::sub_group & sg, float & v0, float & v1, int lane, float inv_sqrt_bs) { + if (ROT > 1) { + { const float u = v0, w = v1; v0 = u + w; v1 = u - w; } // h = 1: pair (2l, 2l+1) + for (int h = 2; h < ROT; h <<= 1) { + const int half = h >> 1; + const float o0 = sycl::permute_group_by_xor(sg, v0, half); + const float o1 = sycl::permute_group_by_xor(sg, v1, half); + const bool upper = (lane & half) != 0; // this lane holds j + h + v0 = upper ? o0 - v0 : v0 + o0; + v1 = upper ? o1 - v1 : v1 + o1; + } + v0 = v0 * inv_sqrt_bs; + v1 = v1 * inv_sqrt_bs; + } +} + +template +void act_row_kernel(const ActParams a, const sycl::nd_item<1> & it, float * red, float * s_rstd) { + const auto sg = it.get_sub_group(); + const int w = (int) sg.get_group_linear_id(); + const int lane = (int) sg.get_local_linear_id(); + const int64_t m = (int64_t) it.get_group(0); + const int64_t K = a.K; + const int nchunks = (int) (K / CHUNK); + const float * x = a.x + m * a.x_stride; + + if (a.gamma) { + if (w == 0) { + float p = 0.0f; + for (int c = lane; c < nchunks; c += SG) + for (int i = 0; i < CHUNK; ++i) { + const float q = x[(int64_t) c * CHUNK + i]; + p = p + (q * q); + } + for (int off = 16; off > 0; off >>= 1) p = p + sycl::permute_group_by_xor(sg, p, off); + if (lane == 0) *s_rstd = 1.0f / sycl::sqrt((p / (float) K) + a.eps); + } + sycl::group_barrier(it.get_group()); + } + const float rstd = a.gamma ? *s_rstd : 0.0f; + + float v0[MAXC], v1[MAXC]; + float mx = 0.0f; + #pragma unroll + for (int i = 0; i < MAXC; ++i) { + const int c = w + i * RW; + if (c < nchunks) { + const int64_t k0 = (int64_t) c * CHUNK + 2 * lane; + float e0 = x[k0], e1 = x[k0 + 1]; + if (a.gamma) { e0 = (e0 * rstd) * a.gamma[k0]; e1 = (e1 * rstd) * a.gamma[k0 + 1]; } + if (a.ascale && a.fold_before) { e0 = e0 / a.ascale[k0]; e1 = e1 / a.ascale[k0 + 1]; } + rot_pair(sg, e0, e1, lane, a.inv_sqrt_bs); + if (a.ascale && !a.fold_before) { e0 = e0 / a.ascale[k0]; e1 = e1 / a.ascale[k0 + 1]; } + v0[i] = e0; v1[i] = e1; + mx = sycl::fmax(mx, sycl::fmax(sycl::fabs(e0), sycl::fabs(e1))); + } + } + for (int off = 16; off > 0; off >>= 1) mx = sycl::fmax(mx, sycl::permute_group_by_xor(sg, mx, off)); + if (lane == 0) red[w] = mx; + sycl::group_barrier(it.get_group()); + mx = red[0]; + for (int j = 1; j < RW; ++j) mx = sycl::fmax(mx, red[j]); + + const float qmax = ABITS == 4 ? 7.0f : 127.0f; + float scale = (a.clip * mx) / qmax; + if (scale < 1e-12f) scale = 1e-12f; + const float inv = 1.0f / scale; + int8_t * row = a.blob + m * a.row_bytes; + #pragma unroll + for (int i = 0; i < MAXC; ++i) { + const int c = w + i * RW; + if (c < nchunks) { + float q0 = sycl::rint(v0[i] * inv), q1 = sycl::rint(v1[i] * inv); + q0 = sycl::fmin(qmax, sycl::fmax(-qmax, q0)); + q1 = sycl::fmin(qmax, sycl::fmax(-qmax, q1)); + if (ABITS == 8) { + row[(int64_t) c * CHUNK + 2 * lane] = (int8_t) (int) q0; + row[(int64_t) c * CHUNK + 2 * lane + 1] = (int8_t) (int) q1; + } else { + row[(int64_t) c * (CHUNK / 2) + lane] = + (int8_t) (((uint32_t) (int) q0 & 0xFu) | (((uint32_t) (int) q1 & 0xFu) << 4)); + } + } + } + if (it.get_local_linear_id() == 0) { + float * sp = (float *) (row + (ABITS == 8 ? K : K / 2)); + *sp = scale; + } +} + +template +void launch_act_row(sycl::queue & q, const ActParams & a) { + q.submit([&](sycl::handler & h) { + sycl::local_accessor red(sycl::range<1>(RW + 1), h); + h.parallel_for(sycl::nd_range<1>((size_t) a.M * RT, RT), [=](sycl::nd_item<1> it) [[sycl::reqd_sub_group_size(SG)]] { + float * r = red.template get_multi_ptr().get(); + act_row_kernel(a, it, r, r + RW); + }); + }); +} + +template +void launch_act_rot(sycl::queue & q, const ActParams & a) { + if (a.K / CHUNK <= RW * MAXC) { + launch_act_row(q, a); + return; + } + const size_t groups = (size_t) ((a.M + ROWS - 1) / ROWS); + q.parallel_for(sycl::nd_range<1>(groups * ROWS * SG, ROWS * SG), + [=](sycl::nd_item<1> it) [[sycl::reqd_sub_group_size(SG)]] { act_kernel(a, it); }); +} + +template +bool launch_act_bits(sycl::queue & q, const ActParams & a, int rot) { + switch (rot) { + case 64: launch_act_rot(q, a); return true; + case 32: launch_act_rot(q, a); return true; + case 16: launch_act_rot(q, a); return true; + case 8: launch_act_rot(q, a); return true; + case 4: launch_act_rot(q, a); return true; + case 2: launch_act_rot(q, a); return true; + case 1: + case 0: launch_act_rot(q, a); return true; + default: return false; + } +} + +// Signed nibble k of a packed row (low nibble = even column). +inline int nib(const int8_t * row, int64_t k) { + const int b = (uint8_t) row[k >> 1]; + const int v = (k & 1) ? (b >> 4) : (b & 0xF); + return v >= 8 ? v - 16 : v; +} + +// One work-item per output element, a TM x TN tile per work-group whose +// activation and weight rows are staged through local memory in K steps. +constexpr int TM = 16, TN = 16, TK = 64; + +template +void launch_gemm_bits(sycl::queue & q, const GemmParams & g, bool has_bias, bool has_res) { + const size_t gm = (size_t) ((g.M + TM - 1) / TM) * TM, gn = (size_t) ((g.N + TN - 1) / TN) * TN; + q.submit([&](sycl::handler & h) { + sycl::local_accessor As(sycl::range<2>(TM, TK), h); + sycl::local_accessor Ws(sycl::range<2>(TN, TK), h); + h.parallel_for(sycl::nd_range<2>(sycl::range<2>(gm, gn), sycl::range<2>(TM, TN)), [=](sycl::nd_item<2> it) { + const int lm = (int) it.get_local_id(0), ln = (int) it.get_local_id(1); + const int64_t m = (int64_t) it.get_global_id(0), n = (int64_t) it.get_global_id(1); + const int64_t m0 = m - lm, n0 = n - ln; + const int tid = lm * TN + ln; + int32_t acc = 0; + for (int64_t k0 = 0; k0 < g.K; k0 += TK) { + // Stage TM activation rows and TN weight rows of TK codes each, unpacked to int8. + for (int e = tid; e < TM * TK; e += TM * TN) { + const int r = e / TK, kk = e % TK; + const int64_t mm = m0 + r; + int8_t v = 0; + if (mm < g.M) { + const int8_t * arow = g.blob + mm * g.row_bytes; + v = ABITS == 4 ? (int8_t) nib(arow, k0 + kk) : arow[k0 + kk]; + } + As[r][kk] = v; + } + for (int e = tid; e < TN * TK; e += TM * TN) { + const int r = e / TK, kk = e % TK; + const int64_t nn = n0 + r; + int8_t v = 0; + if (nn < g.N) { + const int8_t * wrow = g.w + nn * (WBITS == 4 ? g.K / 2 : g.K); + v = WBITS == 4 ? (int8_t) nib(wrow, k0 + kk) : wrow[k0 + kk]; + } + Ws[r][kk] = v; + } + sycl::group_barrier(it.get_group()); + for (int kk = 0; kk < TK; ++kk) + acc += (int32_t) As[lm][kk] * (int32_t) Ws[ln][kk]; + sycl::group_barrier(it.get_group()); + } + if (m < g.M && n < g.N) { + float xs; + std::memcpy(&xs, g.blob + m * g.row_bytes + (ABITS == 4 ? g.K / 2 : g.K), sizeof(float)); + float v = ((float) acc * xs) * g.wscale[n]; + if (has_bias) v = v + g.bias[n]; + if (has_res) v = v + g.res[m * g.N + n]; + g.y[m * g.N + n] = v; + } + }); + }); +} + +// XMX path: int8 joint_matrix (8x32 A, 32x16 B, int32 8x16 accumulator) on +// 16-wide sub-groups. A work-group of XS sub-groups computes a XM x (XS*16) +// tile; per XK-wide K step it stages the activation tile row-major and the +// weight tile in the VNNI layout the B operand wants ([k/4][n*4 + k%4]), +// unpacking INT4 nibbles to int8 on the way, so every width runs on the int8 +// engines. The accumulator is int32 and the epilogue the same as above, so the +// output is still bit-identical to the reference. +namespace jm = sycl::ext::oneapi::experimental::matrix; +constexpr int XM = 32, XS = 4, XN = XS * 16, XK = 32, XT = XS * 16; + +template +void launch_gemm_xmx(sycl::queue & q, const GemmParams & g, bool has_bias, bool has_res) { + const size_t groups_m = (size_t) ((g.M + XM - 1) / XM), groups_n = (size_t) (g.N / XN); + q.submit([&](sycl::handler & h) { + sycl::local_accessor As(sycl::range<1>(XM * XK), h); + sycl::local_accessor Bs(sycl::range<1>(XK * XN), h); + sycl::local_accessor Cs(sycl::range<1>(XM * XN), h); + h.parallel_for(sycl::nd_range<2>(sycl::range<2>(groups_m, groups_n * XT), sycl::range<2>(1, XT)), + [=](sycl::nd_item<2> it) [[sycl::reqd_sub_group_size(16)]] { + const auto sg = it.get_sub_group(); + const int sgid = (int) sg.get_group_linear_id(); + const int tid = (int) it.get_local_id(1); + const int64_t m0 = (int64_t) it.get_group(0) * XM; + const int64_t n0 = (int64_t) it.get_group(1) * XN; + auto a_ptr = As.template get_multi_ptr(); + auto b_ptr = Bs.template get_multi_ptr(); + auto c_ptr = Cs.template get_multi_ptr(); + + jm::joint_matrix acc[XM / 8]; + #pragma unroll + for (int i = 0; i < XM / 8; ++i) jm::joint_matrix_fill(sg, acc[i], 0); + + for (int64_t k0 = 0; k0 < g.K; k0 += XK) { + for (int e = tid; e < XM * XK; e += XT) { + const int r = e / XK, kk = e % XK; + const int64_t mm = m0 + r; + int8_t v = 0; + if (mm < g.M) { + const int8_t * arow = g.blob + mm * g.row_bytes; + v = ABITS == 4 ? (int8_t) nib(arow, k0 + kk) : arow[k0 + kk]; + } + As[e] = v; + } + for (int e = tid; e < XN * XK; e += XT) { + const int nl = e / XK, kk = e % XK; // weight row (output column) nl, element kk + const int8_t * wrow = g.w + (n0 + nl) * (WBITS == 4 ? g.K / 2 : g.K); + const int8_t v = WBITS == 4 ? (int8_t) nib(wrow, k0 + kk) : wrow[k0 + kk]; + Bs[(kk / 4) * (XN * 4) + nl * 4 + (kk % 4)] = v; + } + sycl::group_barrier(it.get_group()); + jm::joint_matrix mb; + jm::joint_matrix_load(sg, mb, b_ptr + sgid * 16 * 4, XN * 4); + #pragma unroll + for (int i = 0; i < XM / 8; ++i) { + jm::joint_matrix ma; + jm::joint_matrix_load(sg, ma, a_ptr + i * 8 * XK, XK); + jm::joint_matrix_mad(sg, acc[i], ma, mb, acc[i]); + } + sycl::group_barrier(it.get_group()); + } + #pragma unroll + for (int i = 0; i < XM / 8; ++i) + jm::joint_matrix_store(sg, acc[i], c_ptr + i * 8 * XN + sgid * 16, XN, jm::layout::row_major); + sycl::group_barrier(it.get_group()); + for (int e = tid; e < XM * XN; e += XT) { + const int r = e / XN, c = e % XN; + const int64_t m = m0 + r, n = n0 + c; + if (m >= g.M) continue; + float xs; + std::memcpy(&xs, g.blob + m * g.row_bytes + (ABITS == 4 ? g.K / 2 : g.K), sizeof(float)); + float v = ((float) Cs[e] * xs) * g.wscale[n]; + if (has_bias) v = v + g.bias[n]; + if (has_res) v = v + g.res[m * g.N + n]; + g.y[m * g.N + n] = v; + } + }); + }); +} + +// Small-M path (the action expert's 10-41 tokens): GEMV-style. A 16-wide +// sub-group owns one output column; its lanes stream the weight row 8 codes at a +// time and accumulate int32 dot products against every token row, which are +// then combined with an exact integer sub-group reduction. Weights are read once +// per column, coalesced across the lanes, which is what a weight-streaming +// shape needs; the epilogue is the reference's. +constexpr int SM_MAX = 32, SM_SG = 16, SM_COLS = 8; // M <= SM_MAX, SM_COLS columns per work-group + +template +inline void load8(const int8_t * p, int64_t k, int8_t * v) { // codes k .. k+7 (k % 8 == 0) + if (BITS == 8) { + for (int j = 0; j < 8; ++j) v[j] = p[k + j]; + } else { + const uint8_t * b = (const uint8_t *) p + (k >> 1); + for (int j = 0; j < 4; ++j) { + const int lo = b[j] & 0xF, hi = b[j] >> 4; + v[2 * j] = (int8_t) (lo >= 8 ? lo - 16 : lo); + v[2 * j + 1] = (int8_t) (hi >= 8 ? hi - 16 : hi); + } + } +} + +template +void launch_gemm_small(sycl::queue & q, const GemmParams & g, bool has_bias, bool has_res) { + const size_t groups = (size_t) ((g.N + SM_COLS - 1) / SM_COLS); + q.parallel_for(sycl::nd_range<1>(groups * SM_COLS * SM_SG, SM_COLS * SM_SG), + [=](sycl::nd_item<1> it) [[sycl::reqd_sub_group_size(SM_SG)]] { + const auto sg = it.get_sub_group(); + const int lane = (int) sg.get_local_linear_id(); + const int64_t n = (int64_t) it.get_group(0) * SM_COLS + (int64_t) sg.get_group_linear_id(); + if (n >= g.N) return; + const int8_t * wrow = g.w + n * (WBITS == 4 ? g.K / 2 : g.K); + int32_t acc[SM_MAX]; + for (int m = 0; m < SM_MAX; ++m) acc[m] = 0; + for (int64_t k = (int64_t) lane * 8; k < g.K; k += SM_SG * 8) { + int8_t wv[8]; + load8(wrow, k, wv); + for (int m = 0; m < SM_MAX; ++m) { + if (m >= g.M) break; + int8_t av[8]; + load8(g.blob + (int64_t) m * g.row_bytes, k, av); + int32_t d = 0; + for (int j = 0; j < 8; ++j) d += (int32_t) wv[j] * (int32_t) av[j]; + acc[m] += d; + } + } + for (int m = 0; m < SM_MAX; ++m) { + if (m >= g.M) break; + const int32_t total = sycl::reduce_over_group(sg, acc[m], sycl::plus()); + if (lane == 0) { + float xs; + std::memcpy(&xs, g.blob + (int64_t) m * g.row_bytes + (ABITS == 4 ? g.K / 2 : g.K), sizeof(float)); + float v = ((float) total * xs) * g.wscale[n]; + if (has_bias) v = v + g.bias[n]; + if (has_res) v = v + g.res[(int64_t) m * g.N + n]; + g.y[(int64_t) m * g.N + n] = v; + } + } + }); +} + +// Whether the device has the int8 8x16x32 XMX combination the kernel above uses. +bool use_xmx(const sycl::queue & q) { + static const bool ok = [&] { + namespace syclex = sycl::ext::oneapi::experimental; + try { + for (const auto & c : q.get_device().get_info()) + if (c.atype == syclex::matrix::matrix_type::sint8 && c.btype == syclex::matrix::matrix_type::sint8 && + c.ctype == syclex::matrix::matrix_type::sint32 && c.nsize == 16 && c.ksize == 32 && + (c.msize == 8 || c.max_msize >= 8)) + return true; + } catch (...) { + } + return false; + }(); + return ok; +} + +bool run_act(sycl::queue & q, ggml_tensor * dst, const FqActSpec & s) { + const ggml_tensor * x = dst->src[0]; + int si = 1; + const ggml_tensor * gamma = s.has_gamma ? dst->src[si++] : nullptr; + const ggml_tensor * ascale = s.has_ascale ? dst->src[si++] : nullptr; + if (!x || x->type != GGML_TYPE_F32 || x->nb[0] != sizeof(float) || s.K % CHUNK != 0 || s.rot_block > CHUNK) + return false; + ActParams a{}; + a.x = (const float *) x->data; + a.x_stride = (int64_t) (x->nb[1] / sizeof(float)); + a.ascale = ascale ? (const float *) ascale->data : nullptr; + a.gamma = gamma ? (const float *) gamma->data : nullptr; + a.blob = (int8_t *) dst->data; + a.row_bytes = dst->ne[0]; + a.M = dst->ne[1] * dst->ne[2] * dst->ne[3]; + a.K = s.K; + a.fold_before = s.fold_before; + a.clip = s.clip; + a.eps = s.eps; + a.inv_sqrt_bs = s.rot_block > 1 ? fqref::inv_sqrt_block(s.rot_block) : 1.0f; + if (a.M == 0) return true; + return s.abits == 8 ? launch_act_bits<8>(q, a, s.rot_block) + : s.abits == 4 ? launch_act_bits<4>(q, a, s.rot_block) + : false; +} + +#ifdef VLA_FQ_SYCL_DNNL +// Prefill shapes on oneDNN's int8 matmul (XMX): INT8 (or unpacked INT4) codes +// in, int32 sums out into a scratch buffer, then the reference epilogue. Integer +// sums are exact in any order, so the output is the same as on the other paths; +// test_foldquant_sycl_op checks that at full K. +struct DnnlState { + dnnl::engine eng; + dnnl::stream strm; + std::map, dnnl::matmul> prims; // M, N, K, lda, s4 weights + int8_t * a8 = nullptr; size_t a8_cap = 0; // A4 activations unpacked to int8 + int8_t * w8 = nullptr; size_t w8_cap = 0; // W4 weights unpacked to int8, when oneDNN takes no s4 + int32_t * acc = nullptr; size_t acc_cap = 0; // the int32 sums + bool s4_ok = true; +}; + +template +T * grow(sycl::queue & q, T *& p, size_t & cap, size_t n) { + if (n > cap) { + if (p) { q.wait(); sycl::free(p, q); } + p = sycl::malloc_device(n, q); + cap = n; + } + return p; +} + +DnnlState & dnnl_state(sycl::queue & q) { + static std::unordered_map states; + auto it = states.find(&q); + if (it == states.end()) { + DnnlState st; + st.eng = dnnl::sycl_interop::make_engine(q.get_device(), q.get_context()); + st.strm = dnnl::sycl_interop::make_stream(st.eng, q); + it = states.emplace(&q, std::move(st)).first; + } + return it->second; +} + +// Unpacks `rows` nibble-packed rows (stride `src_stride` bytes) of K codes to int8. +void unpack_rows(sycl::queue & q, const int8_t * src, int64_t src_stride, int8_t * dst, int64_t rows, int64_t K) { + const int64_t half = K / 2; + q.parallel_for(sycl::range<2>((size_t) rows, (size_t) half), [=](sycl::item<2> it) { + const int64_t r = (int64_t) it.get_id(0), j = (int64_t) it.get_id(1); + const int b = (uint8_t) src[r * src_stride + j]; + const int lo = b & 0xF, hi = b >> 4; + dst[r * K + 2 * j] = (int8_t) (lo >= 8 ? lo - 16 : lo); + dst[r * K + 2 * j + 1] = (int8_t) (hi >= 8 ? hi - 16 : hi); + }); +} + +template +void launch_epilogue(sycl::queue & q, const GemmParams & g, const int32_t * acc, bool has_bias, bool has_res) { + q.parallel_for(sycl::range<2>((size_t) g.M, (size_t) g.N), [=](sycl::item<2> it) { + const int64_t m = (int64_t) it.get_id(0), n = (int64_t) it.get_id(1); + float xs; + std::memcpy(&xs, g.blob + m * g.row_bytes + (ABITS == 4 ? g.K / 2 : g.K), sizeof(float)); + float v = ((float) acc[m * g.N + n] * xs) * g.wscale[n]; + if (has_bias) v = v + g.bias[n]; + if (has_res) v = v + g.res[m * g.N + n]; + g.y[m * g.N + n] = v; + }); +} + +bool gemm_dnnl_impl(sycl::queue & q, const GemmParams & g, int wbits, int abits, bool has_bias, bool has_res); + +// oneDNN on a GPU, unless it cannot drive it (no engine, no int8 matmul): then +// this file's kernels, for good on that queue. GPUs only: on a CPU device +// oneDNN's int8 matmul is not exact everywhere (on CPUs without VNNI it sums +// pairs in saturating 16-bit lanes, vpmaddubsw), and vla.cpp only runs ggml-sycl +// on GPUs anyway; the CPU device is test_foldquant_sycl_op's stand-in. +bool gemm_dnnl(sycl::queue & q, const GemmParams & g, int wbits, int abits, bool has_bias, bool has_res) { + static std::unordered_map unusable; + if (unusable.find(&q) == unusable.end()) unusable[&q] = !q.get_device().is_gpu(); + if (unusable[&q]) return false; + try { + if (gemm_dnnl_impl(q, g, wbits, abits, has_bias, has_res)) return true; + } catch (const std::exception & e) { + std::fprintf(stderr, "FoldQuant SYCL: oneDNN unavailable on this device (%s); using the native GEMMs\n", e.what()); + } + unusable[&q] = true; + return false; +} + +bool gemm_dnnl_impl(sycl::queue & q, const GemmParams & g, int wbits, int abits, bool has_bias, bool has_res) { + using dt = dnnl::memory::data_type; + using md = dnnl::memory::desc; + DnnlState & st = dnnl_state(q); + const int64_t M = g.M, N = g.N, K = g.K; + + const int8_t * a = g.blob; + int64_t lda = g.row_bytes; + if (abits == 4) { + unpack_rows(q, g.blob, g.row_bytes, grow(q, st.a8, st.a8_cap, (size_t) (M * K)), M, K); + a = st.a8; + lda = K; + } + bool s4 = wbits == 4 && st.s4_ok; + auto make = [&](bool w4) { + const md src_md({ M, K }, dt::s8, { lda, 1 }); + const md wei_md({ K, N }, w4 ? dt::s4 : dt::s8, { 1, K }); + const md dst_md({ M, N }, dt::s32, { N, 1 }); + return dnnl::matmul(dnnl::matmul::primitive_desc(st.eng, src_md, wei_md, dst_md)); + }; + auto key = std::make_tuple(M, N, K, lda, s4 ? 1 : 0); + auto it = st.prims.find(key); + if (it == st.prims.end()) { + try { + it = st.prims.emplace(key, make(s4)).first; + } catch (const dnnl::error &) { + if (!s4) return false; + st.s4_ok = false; // no s4 x s8 matmul here: unpack the weights instead + std::fprintf(stderr, "FoldQuant SYCL: oneDNN has no s8 x s4 matmul here; INT4 weights are unpacked per call\n"); + s4 = false; + key = std::make_tuple(M, N, K, lda, 0); + it = st.prims.find(key); + if (it == st.prims.end()) it = st.prims.emplace(key, make(false)).first; + } + } + const int8_t * w = g.w; + if (wbits == 4 && !s4) { + unpack_rows(q, g.w, K / 2, grow(q, st.w8, st.w8_cap, (size_t) (N * K)), N, K); + w = st.w8; + } + int32_t * acc = grow(q, st.acc, st.acc_cap, (size_t) (M * N)); + using dnnl::sycl_interop::make_memory; + using dnnl::sycl_interop::memory_kind; + const md src_md({ M, K }, dt::s8, { lda, 1 }); + const md wei_md({ K, N }, s4 ? dt::s4 : dt::s8, { 1, K }); + const md dst_md({ M, N }, dt::s32, { N, 1 }); + auto src_m = make_memory(src_md, st.eng, memory_kind::usm, const_cast(a)); + auto wei_m = make_memory(wei_md, st.eng, memory_kind::usm, const_cast(w)); + auto dst_m = make_memory(dst_md, st.eng, memory_kind::usm, acc); + it->second.execute(st.strm, { { DNNL_ARG_SRC, src_m }, { DNNL_ARG_WEIGHTS, wei_m }, { DNNL_ARG_DST, dst_m } }); + if (abits == 4) launch_epilogue<4>(q, g, acc, has_bias, has_res); + else launch_epilogue<8>(q, g, acc, has_bias, has_res); + return true; +} +#endif + +bool run_gemm(sycl::queue & q, ggml_tensor * dst, const FqGemmSpec & s, int abits) { + const ggml_tensor * w = dst->src[0], * blob = dst->src[1], * ws = dst->src[2]; + const ggml_tensor * bias = dst->src[3], * res = dst->src[4]; + if (!w || !blob || !ws || s.heads != 0) return false; + GemmParams g{}; + g.w = (const int8_t *) w->data; + g.blob = (const int8_t *) blob->data; + g.wscale = (const float *) ws->data; + g.bias = bias ? (const float *) bias->data : nullptr; + g.res = res ? (const float *) res->data : nullptr; + g.y = (float *) dst->data; + g.M = dst->ne[1]; + g.N = s.N; + g.K = s.K; + g.row_bytes = blob->ne[0]; + if (g.M == 0) return true; + // oneDNN when built with it; VLA_FQ_SYCL_GEMM=native picks this file's kernels + // (the GEMV for M <= SM_MAX, else XMX) and =simple the plain tiled one. Read + // per call so test_foldquant_sycl_op can cover every path. + const char * mode = std::getenv("VLA_FQ_SYCL_GEMM"); + const bool simple = mode && std::strcmp(mode, "simple") == 0; + const bool native = simple || (mode && std::strcmp(mode, "native") == 0); +#ifdef VLA_FQ_SYCL_DNNL + if (!native && gemm_dnnl(q, g, s.wbits, abits, bias, res)) + return true; +#endif + if (!simple && g.M <= SM_MAX && g.K % (SM_SG * 8) == 0) { + if (s.wbits == 8 && abits == 8) launch_gemm_small<8, 8>(q, g, bias, res); + else if (s.wbits == 8 && abits == 4) launch_gemm_small<8, 4>(q, g, bias, res); + else if (s.wbits == 4 && abits == 8) launch_gemm_small<4, 8>(q, g, bias, res); + else if (s.wbits == 4 && abits == 4) launch_gemm_small<4, 4>(q, g, bias, res); + else return false; + return true; + } + if (!simple && use_xmx(q) && g.N % XN == 0 && g.K % XK == 0) { + if (s.wbits == 8 && abits == 8) launch_gemm_xmx<8, 8>(q, g, bias, res); + else if (s.wbits == 8 && abits == 4) launch_gemm_xmx<8, 4>(q, g, bias, res); + else if (s.wbits == 4 && abits == 8) launch_gemm_xmx<4, 8>(q, g, bias, res); + else if (s.wbits == 4 && abits == 4) launch_gemm_xmx<4, 4>(q, g, bias, res); + else return false; + return true; + } + if (s.wbits == 8 && abits == 8) launch_gemm_bits<8, 8>(q, g, bias, res); + else if (s.wbits == 8 && abits == 4) launch_gemm_bits<8, 4>(q, g, bias, res); + else if (s.wbits == 4 && abits == 8) launch_gemm_bits<4, 8>(q, g, bias, res); + else if (s.wbits == 4 && abits == 4) launch_gemm_bits<4, 4>(q, g, bias, res); + else return false; + return true; +} + +// The gemm node's activation width is the one its blob was quantized at: the +// blob row holds K (A8) or K/2 (A4) code bytes then the scale and its padding. +int blob_abits(const ggml_tensor * blob, int64_t K) { + return blob->ne[0] == fq_act_row_bytes(K, 4) ? 4 : 8; +} + +// VLA_FQ_CHECK=1, as on CUDA: after each node, recompute it with the CPU +// reference on host copies of the same inputs and report any byte difference. +// Slow; a diagnostic for a mismatch the unit tests' shapes do not reproduce. +void check_against_host(sycl::queue & q, ggml_tensor * dst, void (*ref)(ggml_tensor *, int, int, void *), void * ud, + size_t cmp_bytes_per_row, size_t row_bytes) { + std::vector src_bytes[GGML_MAX_SRC]; + ggml_tensor src_copy[GGML_MAX_SRC]; + ggml_tensor d = *dst; + q.wait(); + for (int i = 0; i < GGML_MAX_SRC; ++i) { + if (!dst->src[i]) { d.src[i] = nullptr; continue; } + src_copy[i] = *dst->src[i]; + src_bytes[i].resize(ggml_nbytes(dst->src[i])); + q.memcpy(src_bytes[i].data(), dst->src[i]->data, src_bytes[i].size()); + src_copy[i].data = src_bytes[i].data(); + d.src[i] = &src_copy[i]; + } + std::vector got(ggml_nbytes(dst)), want(ggml_nbytes(dst)); + q.memcpy(got.data(), dst->data, got.size()); + q.wait(); + d.data = want.data(); + ref(&d, 0, 1, ud); + const int64_t rows = ggml_nbytes(dst) / row_bytes; + int64_t bad = 0, first = -1; + for (int64_t r = 0; r < rows; ++r) + if (std::memcmp(got.data() + (size_t) r * row_bytes, want.data() + (size_t) r * row_bytes, cmp_bytes_per_row)) { + if (first < 0) first = r; + ++bad; + } + std::printf("vla(fq) CHECK %-40s ne=[%lld,%lld] %s", ggml_get_name(dst), (long long) dst->ne[0], + (long long) dst->ne[1], bad ? "MISMATCH" : "ok\n"); + if (bad) std::printf(" rows %lld/%lld (first %lld)\n", (long long) bad, (long long) rows, (long long) first); +} + +bool dispatch(ggml_tensor * dst, void * queue) { + const void * ud = fq_userdata(dst); + const uint32_t m = fq_magic(ud); + if (m != FQ_ACT_MAGIC && m != FQ_GEMM_MAGIC) return false; + sycl::queue & q = *static_cast(queue); + static const bool check = env_flag("VLA_FQ_CHECK"); + if (m == FQ_ACT_MAGIC) { + const FqActSpec & s = *static_cast(ud); + if (!run_act(q, dst, s)) return false; + if (check) + check_against_host(q, dst, fq_act_cpu, const_cast(ud), (size_t) fq_act_kpack(s.K, s.abits) + 4, + (size_t) dst->ne[0]); + return true; + } + const FqGemmSpec & s = *static_cast(ud); + if (!run_gemm(q, dst, s, blob_abits(dst->src[1], s.K))) return false; + if (check) + check_against_host(q, dst, fq_gemm_cpu, const_cast(ud), (size_t) s.N * sizeof(float), + (size_t) s.N * sizeof(float)); + return true; +} + +bool supports(const ggml_tensor * op) { + const uint32_t m = fq_magic(fq_userdata(op)); + if (m == FQ_ACT_MAGIC) return true; + if (m == FQ_GEMM_MAGIC) return static_cast(fq_userdata(op))->heads == 0; + return false; +} + +} // namespace + +void sycl_register_foldquant_ops() { + ggml_sycl_ext_forward = dispatch; + ggml_sycl_ext_supports = supports; +} + +} // namespace vla diff --git a/src/sycl/vla_sycl_foldquant.h b/src/sycl/vla_sycl_foldquant.h new file mode 100644 index 0000000..dd0d458 --- /dev/null +++ b/src/sycl/vla_sycl_foldquant.h @@ -0,0 +1,25 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant on ggml's SYCL backend (src/sycl/vla_sycl_foldquant.cpp). + +#pragma once + +namespace vla { + +// Points the SYCL extension hook (scripts/patch_ggml_sycl_ext_hook.py) at the +// FoldQuant kernels. Idempotent; call once the SYCL backend is up. +void sycl_register_foldquant_ops(); + +} // namespace vla diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 545955f..d8be335 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -61,6 +61,43 @@ target_link_libraries(test_prompt PRIVATE vla_core) target_compile_options(test_prompt PRIVATE -Wall -Wextra) add_test(NAME prompt COMMAND test_prompt) +# FoldQuant CPU reference ops as GGML_OP_CUSTOM nodes on the CPU backend. The +# golden checksum is shared with tests/py/test_foldquant_ref.py. +add_executable(test_foldquant_cpu_op test_foldquant_cpu_op.cpp) +target_include_directories(test_foldquant_cpu_op PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_link_libraries(test_foldquant_cpu_op PRIVATE vla_core) +target_compile_options(test_foldquant_cpu_op PRIVATE -Wall -Wextra -ffp-contract=off) +add_test(NAME foldquant_cpu_op COMMAND test_foldquant_cpu_op) + +# FoldQuant sites read back as float weights (the path every backend without +# the integer kernels takes) against the dense rotation they stand for. +add_executable(test_foldquant_dequant test_foldquant_dequant.cpp) +target_include_directories(test_foldquant_dequant PRIVATE ${CMAKE_SOURCE_DIR}/src) +target_link_libraries(test_foldquant_dequant PRIVATE vla_core) +target_compile_options(test_foldquant_dequant PRIVATE -Wall -Wextra) +add_test(NAME foldquant_dequant COMMAND test_foldquant_dequant) + +# The same site through ggml's SYCL backend (src/sycl/vla_sycl_foldquant.cpp), +# bit-identical to the CPU reference. +if(GGML_SYCL) + add_executable(test_foldquant_sycl_op test_foldquant_sycl_op.cpp) + target_include_directories(test_foldquant_sycl_op PRIVATE ${CMAKE_SOURCE_DIR}/src) + target_link_libraries(test_foldquant_sycl_op PRIVATE vla_core) + # It runs the CPU reference inline (no contraction) and, without a GPU, + # drives the kernels on a SYCL queue of its own. + target_compile_options(test_foldquant_sycl_op PRIVATE -fsycl -Wall -Wextra -ffp-contract=off) + add_test(NAME foldquant_sycl_op COMMAND test_foldquant_sycl_op) +endif() + +# The same site through ggml's OpenVINO backend (src/openvino/foldquant_ov.cpp). +if(GGML_OPENVINO) + add_executable(test_foldquant_ov_op test_foldquant_ov_op.cpp) + target_include_directories(test_foldquant_ov_op PRIVATE ${CMAKE_SOURCE_DIR}/src) + target_link_libraries(test_foldquant_ov_op PRIVATE vla_core) + target_compile_options(test_foldquant_ov_op PRIVATE -Wall -Wextra) + add_test(NAME foldquant_ov_op COMMAND test_foldquant_ov_op) +endif() + # A/B harness for the two BitVLA ternary-GEMM tilings. Built so it cannot rot, # not registered with ctest: it needs a GPU and is read by hand. # VLA_BITVLA_NARROW_GEMM=1 selects the old one-tile-per-CTA kernel at runtime. @@ -70,12 +107,33 @@ if(GGML_CUDA) target_link_libraries(bitvla_gemm_check PRIVATE bitvla_cuda_kernels) endif() +# FoldQuant kernel micro-benchmark at GR00T shapes against cuBLAS BF16. Built +# so it cannot rot, not registered with ctest: it needs a GPU and is read by hand. +if(GGML_CUDA) + add_executable(foldquant_gemm_check foldquant_gemm_check.cu) + target_include_directories(foldquant_gemm_check PRIVATE ${CMAKE_SOURCE_DIR}/src) + target_link_libraries(foldquant_gemm_check PRIVATE vla_fq_kernels CUDA::cublas) +endif() + +# FoldQuant CUDA kernels versus the CPU reference, byte-identical; run twice, +# the second time through the VLA_FQ_CPU_REF=1 staging shim. Skips (exit 0) +# without a CUDA device. +if(GGML_CUDA) + add_executable(test_foldquant_cuda_op test_foldquant_cuda_op.cpp) + target_include_directories(test_foldquant_cuda_op PRIVATE ${CMAKE_SOURCE_DIR}/src) + target_link_libraries(test_foldquant_cuda_op PRIVATE vla_core vla_cuda_ops ggml) + target_compile_options(test_foldquant_cuda_op PRIVATE -Wall -Wextra -ffp-contract=off) + add_test(NAME foldquant_cuda_op COMMAND test_foldquant_cuda_op) + add_test(NAME foldquant_cuda_op_ref COMMAND test_foldquant_cuda_op) + set_tests_properties(foldquant_cuda_op_ref PROPERTIES ENVIRONMENT "VLA_FQ_CPU_REF=1;GGML_CUDA_DISABLE_GRAPHS=1") +endif() + # Regression test for the in-tree BF16 CUDA kernels and the ggml hook they ride # on. Skips itself (exit 0) when no CUDA device is present. if(GGML_CUDA) add_executable(test_bf16_cuda_ops test_bf16_cuda_ops.cpp) target_include_directories(test_bf16_cuda_ops PRIVATE ${CMAKE_SOURCE_DIR}/src) - target_link_libraries(test_bf16_cuda_ops PRIVATE vla_cuda_ops ggml) + target_link_libraries(test_bf16_cuda_ops PRIVATE vla_core vla_cuda_ops ggml) target_compile_options(test_bf16_cuda_ops PRIVATE -Wall -Wextra) add_test(NAME bf16_cuda_ops COMMAND test_bf16_cuda_ops) add_test(NAME bf16_cuda_ops_flat COMMAND test_bf16_cuda_ops) diff --git a/tests/foldquant_gemm_check.cu b/tests/foldquant_gemm_check.cu new file mode 100644 index 0000000..c498fde --- /dev/null +++ b/tests/foldquant_gemm_check.cu @@ -0,0 +1,194 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Micro-benchmark of the FoldQuant CUDA kernels at GR00T N1.6/N1.7 shapes: +// prologue + INT8 GEMM time, effective weight-streaming bandwidth, and a +// cuBLAS BF16 GEMM of the same shape as the number to beat. Built with the +// CUDA tree, not registered with ctest (needs a GPU, read by hand). +// +// ./build-cuda/tests/foldquant_gemm_check [reps] + +#include "kernels/foldquant/fq_kernels.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#define CK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { std::printf("CUDA error %s at %d\n", cudaGetErrorString(e), __LINE__); return 1; } } while (0) + +struct Shape { const char * name; int M, N, K; bool gamma; }; + +// --stress: run the FFN shapes many times and hash the outputs; any change +// between iterations is a race in the kernels (the graph is otherwise fixed). +static uint64_t fnv(const void * p, size_t n) { + uint64_t h = 14695981039346656037ull; + for (size_t i = 0; i < n; ++i) { h ^= ((const uint8_t *) p)[i]; h *= 1099511628211ull; } + return h; +} + +static int stress(cudaStream_t st, int iters) { + const Shape shapes[] = { + { "ffn gate M=87", 87, 6144, 2048, true }, + { "ffn down M=87", 87, 2048, 6144, false }, + { "qkv M=87", 87, 2048, 2048, true }, + }; + for (const Shape & s : shapes) { + const int64_t M = s.M, N = s.N, K = s.K, rb = K + 16; + std::vector hx((size_t) M * K), hg(K), hws(N); + std::vector hw((size_t) N * K); + for (size_t i = 0; i < hx.size(); ++i) hx[i] = (float) ((i * 37) % 23) - 11.0f; + for (size_t i = 0; i < hw.size(); ++i) hw[i] = (int8_t) ((i * 53) % 255 - 127); + for (int64_t i = 0; i < K; ++i) hg[i] = 1.0f + 0.001f * (float) (i % 7); + for (int64_t i = 0; i < N; ++i) hws[i] = 0.01f; + float *x, *g, *ws, *y; int8_t *w, *blob; + CK(cudaMalloc(&x, hx.size() * 4)); CK(cudaMalloc(&g, K * 4)); CK(cudaMalloc(&ws, N * 4)); + CK(cudaMalloc(&y, (size_t) M * N * 4)); CK(cudaMalloc(&w, hw.size())); CK(cudaMalloc(&blob, (size_t) M * rb)); + CK(cudaMemcpy(x, hx.data(), hx.size() * 4, cudaMemcpyHostToDevice)); + CK(cudaMemcpy(w, hw.data(), hw.size(), cudaMemcpyHostToDevice)); + CK(cudaMemcpy(g, hg.data(), K * 4, cudaMemcpyHostToDevice)); + CK(cudaMemcpy(ws, hws.data(), N * 4, cudaMemcpyHostToDevice)); + vla::fq::ActArgs a{}; a.x = x; a.gamma = s.gamma ? g : nullptr; a.blob = blob; a.row_bytes = rb; + a.M = M; a.K = K; a.abits = 8; a.rot_block = 64; a.clip = 1.0f; a.eps = 1e-6f; a.inv_sqrt_bs = 0.125f; + vla::fq::GemmArgs gm{}; gm.w = w; gm.blob = blob; gm.wscale = ws; gm.y = y; gm.M = M; gm.N = N; gm.K = K; + gm.row_bytes = rb; gm.wbits = 8; gm.abits = 8; + std::vector hb((size_t) M * rb); std::vector hy((size_t) M * N); + uint64_t hb0 = 0, hy0 = 0; int bad_b = 0, bad_y = 0; + for (int it = 0; it < iters; ++it) { + CK(cudaMemsetAsync(blob, 0xAB, (size_t) M * rb, st)); + CK(cudaMemsetAsync(y, 0xCD, (size_t) M * N * 4, st)); + if (vla::fq::launch_act(a, st) != cudaSuccess) { std::printf("act launch failed\n"); return 1; } + if (vla::fq::launch_gemm(gm, st) != cudaSuccess) { std::printf("gemm launch failed\n"); return 1; } + CK(cudaMemcpyAsync(hb.data(), blob, hb.size(), cudaMemcpyDeviceToHost, st)); + CK(cudaMemcpyAsync(hy.data(), y, hy.size() * 4, cudaMemcpyDeviceToHost, st)); + CK(cudaStreamSynchronize(st)); + // hash codes + scale of each row (padding bytes are never written) + uint64_t hbb = 14695981039346656037ull; + for (int64_t m = 0; m < M; ++m) hbb ^= fnv(hb.data() + m * rb, K + 4), hbb *= 1099511628211ull; + const uint64_t hyy = fnv(hy.data(), hy.size() * 4); + if (it == 0) { hb0 = hbb; hy0 = hyy; } + else { bad_b += hbb != hb0; bad_y += hyy != hy0; } + } + std::printf("stress %-16s blob changed %d/%d y changed %d/%d\n", s.name, bad_b, iters - 1, bad_y, iters - 1); + cudaFree(x); cudaFree(g); cudaFree(ws); cudaFree(y); cudaFree(w); cudaFree(blob); + } + return 0; +} + +int main(int argc, char ** argv) { + if (argc > 1 && std::string(argv[1]) == "--stress") { + cudaStream_t st; CK(cudaStreamCreate(&st)); + return stress(st, argc > 2 ? std::atoi(argv[2]) : 200); + } + // --cold: rotate every kernel over enough weight copies that no site stays + // in the 4 MB L2 between reps - what a model run sees, where each site's + // weights are streamed from DRAM once per request. Without it the loop + // re-runs one L2-hot site and overstates the achievable bandwidth. + bool cold = false; + int reps = 50; + for (int i = 1; i < argc; ++i) { + if (std::string(argv[i]) == "--cold") cold = true; + else reps = std::atoi(argv[i]); + } + const Shape shapes[] = { + { "dit qkv M=41", 41, 4608, 1536, false }, + { "dit o M=41", 41, 1536, 1536, false }, + { "dit ff0 M=41", 41, 6144, 1536, false }, + { "dit ff2 M=41", 41, 1536, 6144, false }, + { "llm qkv M=160", 160, 4096, 2048, true }, + { "llm o M=160", 160, 2048, 2048, false }, + { "llm gate+up M=160", 160, 12288, 2048, true }, + { "llm down M=160", 160, 2048, 6144, false }, + { "llm qkv M=1024", 1024, 4096, 2048, true }, + }; + cublasHandle_t cublas; cublasCreate(&cublas); + cudaStream_t st; CK(cudaStreamCreate(&st)); cublasSetStream(cublas, st); + cudaEvent_t e0, e1; CK(cudaEventCreate(&e0)); CK(cudaEventCreate(&e1)); + + { // burn-in so the governor raises the clock before anything is timed + float * junk; CK(cudaMalloc(&junk, 64 << 20)); + for (int i = 0; i < 200; ++i) CK(cudaMemsetAsync(junk, i, 64 << 20, st)); + CK(cudaStreamSynchronize(st)); cudaFree(junk); + } + std::printf("%-20s %8s %8s %8s %9s %9s %9s\n", "shape", "prolog us", "gemm us", "total us", "W GB/s", "bf16 us", "speedup"); + for (const Shape & s : shapes) { + const int64_t M = s.M, N = s.N, K = s.K; + const int64_t rb = K + 16; + std::vector hx((size_t) M * K), hg(K), hws(N); + std::vector hw((size_t) N * K); + for (size_t i = 0; i < hx.size(); ++i) hx[i] = (float) ((i * 37) % 23) - 11.0f; + for (size_t i = 0; i < hw.size(); ++i) hw[i] = (int8_t) ((i * 53) % 255 - 127); + for (int64_t i = 0; i < K; ++i) hg[i] = 1.0f; + for (int64_t i = 0; i < N; ++i) hws[i] = 0.01f; + + const int copies = cold ? (int) std::max(2, (48ll << 20) / ((int64_t) N * K) + 1) : 1; + float *x, *g, *ws, *y; int8_t *blob; + std::vector w(copies); + CK(cudaMalloc(&x, hx.size() * 4)); CK(cudaMalloc(&g, K * 4)); CK(cudaMalloc(&ws, N * 4)); + CK(cudaMalloc(&y, (size_t) M * N * 4)); CK(cudaMalloc(&blob, (size_t) M * rb)); + for (int c = 0; c < copies; ++c) { + CK(cudaMalloc(&w[c], hw.size())); + CK(cudaMemcpy(w[c], hw.data(), hw.size(), cudaMemcpyHostToDevice)); + } + int rot = 0; + CK(cudaMemcpy(x, hx.data(), hx.size() * 4, cudaMemcpyHostToDevice)); + CK(cudaMemcpy(g, hg.data(), K * 4, cudaMemcpyHostToDevice)); + CK(cudaMemcpy(ws, hws.data(), N * 4, cudaMemcpyHostToDevice)); + + vla::fq::ActArgs a{}; a.x = x; a.ascale = nullptr; a.gamma = s.gamma ? g : nullptr; a.blob = blob; a.row_bytes = rb; + a.M = M; a.K = K; a.abits = 8; a.rot_block = 64; a.fold_before = false; a.clip = 1.0f; a.eps = 1e-6f; a.inv_sqrt_bs = 0.125f; + vla::fq::GemmArgs gm{}; gm.w = w[0]; gm.blob = blob; gm.wscale = ws; gm.bias = nullptr; gm.y = y; gm.M = M; gm.N = N; gm.K = K; + gm.row_bytes = rb; gm.wbits = 8; gm.abits = 8; + + // DVFS: the governor ramps the clock with sustained load, so the three + // kernels are measured interleaved in several rounds and the best round + // of each is kept - the ratios are what matters, at whatever clock. + auto time = [&](auto fn) -> float { + for (int i = 0; i < 5; ++i) fn(); + CK(cudaEventRecord(e0, st)); + for (int i = 0; i < reps; ++i) fn(); + CK(cudaEventRecord(e1, st)); CK(cudaEventSynchronize(e1)); + float ms; cudaEventElapsedTime(&ms, e0, e1); return ms * 1000.0f / reps; + }; + std::vector<__nv_bfloat16 *> wb(copies); __nv_bfloat16 * xb; float * yb; + for (int c = 0; c < copies; ++c) CK(cudaMalloc(&wb[c], (size_t) N * K * 2)); + CK(cudaMalloc(&xb, (size_t) M * K * 2)); CK(cudaMalloc(&yb, (size_t) M * N * 4)); + const float one = 1.0f, zero = 0.0f; + auto bf16 = [&] { + // cuBLAS BF16 GEMM with F32 accumulate, the shape vla.cpp's bf16 path runs (weights [N,K], x [M,K]). + cublasGemmEx(cublas, CUBLAS_OP_T, CUBLAS_OP_N, (int) N, (int) M, (int) K, &one, wb[rot++ % copies], CUDA_R_16BF, (int) K, + xb, CUDA_R_16BF, (int) K, &zero, yb, CUDA_R_32F, (int) N, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); + }; + float t_pro = 1e30f, t_gemm = 1e30f, t_bf16 = 1e30f; + for (int round = 0; round < 4; ++round) { + t_bf16 = std::min(t_bf16, time(bf16)); + t_pro = std::min(t_pro, time([&] { vla::fq::launch_act(a, st); })); + t_gemm = std::min(t_gemm, time([&] { gm.w = w[rot++ % copies]; vla::fq::launch_gemm(gm, st); })); + } + if (cudaGetLastError() != cudaSuccess) { std::printf("launch error\n"); return 1; } + const double gbps = (double) N * K / ((t_gemm) * 1e-6) / 1e9; + std::printf("%-20s %8.1f %8.1f %8.1f %9.1f %9.1f %8.2fx\n", s.name, t_pro, t_gemm, t_pro + t_gemm, gbps, t_bf16, + t_bf16 / (t_pro + t_gemm)); + cudaFree(x); cudaFree(g); cudaFree(ws); cudaFree(y); cudaFree(blob); cudaFree(xb); cudaFree(yb); + for (int c = 0; c < copies; ++c) { cudaFree(w[c]); cudaFree(wb[c]); } + } + return 0; +} diff --git a/tests/py/test_foldquant_ref.py b/tests/py/test_foldquant_ref.py new file mode 100644 index 0000000..3829600 --- /dev/null +++ b/tests/py/test_foldquant_ref.py @@ -0,0 +1,93 @@ +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +"""Pins scripts/foldquant_ref.py to the C++ reference (src/foldquant_ref.h): +the golden checksum of the codes for the fixed LCG input is the same constant +tests/test_foldquant_cpu_op.cpp asserts, so the two references cannot drift +apart. Also pins the nibble order, the blob layout and rotation_block_for.""" + +import pathlib +import re +import sys + +import pytest + +np = pytest.importorskip("numpy") + +SCRIPTS = pathlib.Path(__file__).resolve().parents[2] / "scripts" +sys.path.insert(0, str(SCRIPTS)) + +import foldquant_ref as fq # noqa: E402 + +K, T = 128, 5 + + +def _golden_from_cpp() -> int: + src = (pathlib.Path(__file__).resolve().parents[1] / "test_foldquant_cpu_op.cpp").read_text() + m = re.search(r"FQ_GOLDEN_FNV\s*=\s*0x([0-9a-fA-F]+)", src) + assert m, "golden constant missing from test_foldquant_cpu_op.cpp" + return int(m.group(1), 16) + + +def test_golden_checksum_matches_cpp(): + x = fq.lcg(0x5EED1234, K * T).reshape(T, K) * np.float32(4.0) + codes, _ = fq.act_quant(x, bits=8, rot_block=64) + assert fq.fnv1a(codes.tobytes()) == _golden_from_cpp() + + +def test_fwht_matches_dense_hadamard(): + rng = np.random.default_rng(0) + x = rng.standard_normal((3, 256)).astype(np.float32) + h = fq.hadamard(64) + dense = (x.reshape(3, 4, 64).astype(np.float64) @ h.T).reshape(3, 256) + assert np.allclose(fq.fwht_rows(x, 64), dense, atol=1e-5) + + +def test_fold_then_rotate_is_identity_up_to_quant(): + # W' x' with x' = H x and W' = W H^T equals W x (H orthogonal). + rng = np.random.default_rng(1) + w = rng.standard_normal((16, 128)).astype(np.float32) + x = rng.standard_normal((4, 128)).astype(np.float32) + wf = fq.fold_weight(w, 64) + xr = fq.fwht_rows(x, 64) + assert np.allclose(xr @ wf.T, x @ w.T, atol=1e-4) + + +def test_nibble_round_trip_and_order(): + codes = np.array([[-7, 7, 0, -1, 3, -8 + 1, 1, 2]], dtype=np.int8) + packed = fq.pack_nibbles(codes) + assert packed[0, 0] == ((-7 & 0xF) | ((7 & 0xF) << 4)) # low nibble = even column + assert np.array_equal(fq.unpack_nibbles(packed), codes) + + +def test_blob_layout(): + codes = np.zeros((2, 64), dtype=np.int8) + scales = np.array([1.5, -2.0], dtype=np.float32) + blob = fq.act_blob(codes, scales, bits=8) + assert blob.shape == (2, 64 + fq.ACT_TAIL) + assert blob[:, 64:68].copy().view(np.float32).reshape(-1).tolist() == [1.5, -2.0] + blob4 = fq.act_blob(codes, scales, bits=4) + assert blob4.shape == (2, 32 + fq.ACT_TAIL) + + +def test_rotation_block_for(): + assert fq.rotation_block_for(2048, 64) == 64 + assert fq.rotation_block_for(480, 64) == 32 + assert fq.rotation_block_for(100, 64) == 4 + assert fq.rotation_block_for(7, 64) == 1 + + +def test_gemm_ref_epilogue(): + wc = np.array([[1, 2], [3, 4]], dtype=np.int8) + xc = np.array([[5, 6]], dtype=np.int8) + y = fq.gemm_ref(wc, np.array([0.5, 0.25], np.float32), xc, np.array([2.0], np.float32), np.array([1.0, 0.0], np.float32)) + assert y.tolist() == [[(5 + 12) * 2 * 0.5 + 1.0, (15 + 24) * 2 * 0.25]] diff --git a/tests/py/test_quantized_model_converter.py b/tests/py/test_quantized_model_converter.py new file mode 100644 index 0000000..2f18ad4 --- /dev/null +++ b/tests/py/test_quantized_model_converter.py @@ -0,0 +1,473 @@ +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for scripts/convert_quantized_model_to_gguf.py and scripts/gguf_quant_writer.py. + +A tiny GR00T-shaped model is recorded three ways (a FoldQuantVLA quantized +checkpoint, an earlier FoldQuantVLA fake-quant state, a plugin ONNX graph pair) +and every route must produce the same GGUF sites. The end-to-end check against real exports runs through `--check-onnx`. +""" + +import json +import pathlib +import sys +import types + +import numpy as np +import pytest +import torch + +SCRIPTS = pathlib.Path(__file__).resolve().parents[2] / "scripts" +sys.path.insert(0, str(SCRIPTS)) + +# test_converters.py installs stub numpy/torch/gguf/safetensors modules (it runs on bare +# Python in CI, one file per process); collected in the same pytest process, those stubs +# shadow the real packages this test needs. +if not hasattr(sys.modules.get("safetensors"), "__path__") and "safetensors" in sys.modules: + pytest.skip("real safetensors/gguf are stubbed by another test module in this process; " + "run this file on its own", allow_module_level=True) + +import gguf # noqa: E402 +from safetensors.numpy import save_file # noqa: E402 +from safetensors.torch import save_file as save_pt # noqa: E402 + +import convert_quantized_model_to_gguf as cf # noqa: E402 +from gguf_quant_writer import QuantizingGGUFWriter, Site # noqa: E402 + +# A two-layer Qwen3-shaped LLM (K=128, q 128 / kv 64, inner 128) over a two-block DiT +# (K=128, inner 64, encoder 256, ffn 256): block 0 cross-attention, block 1 self-attention. +K, KV, INNER, D, K_ENC, FF = 128, 64, 128, 64, 256, 256 +LM_ROOT, AHK = "backbone.lm", "action_head" +RNG = np.random.default_rng(0) + + +def _i8(n, k): + return RNG.integers(-127, 128, (n, k), dtype=np.int8) + + +def _u8(n, k): + return RNG.integers(0, 256, (n, k // 2), dtype=np.uint8) + + +def _f32(n, lo=0.5, hi=2.0): + return RNG.uniform(lo, hi, n).astype(np.float32) + + +def _bf16(n): + return (torch.randn(n) * 0.1).to(torch.bfloat16) + + +class Tiny: + """The recorded model: LLM sites int8, DiT sites int4, plus the checkpoint gains and SQ vectors.""" + + def __init__(self): + self.llm = {} + for i in range(2): + self.llm[f"L{i}_qkv"] = (_i8(K + 2 * KV, K), _f32(K + 2 * KV)) + self.llm[f"L{i}_o"] = (_i8(K, K), _f32(K)) + self.llm[f"L{i}_gateup"] = (_i8(2 * INNER, K), _f32(2 * INNER)) + self.llm[f"L{i}_down"] = (_i8(K, INNER), _f32(K)) + self.sq = {f"L{i}_{s}": _f32(INNER if s == "down" else K) for i in range(2) for s in ("qkv", "gateup", "down")} + self.gains = {f"L{i}_{n}": _bf16(K) for i in range(2) for n in ("input_layernorm", "post_attention_layernorm")} + self.dit = { + "block0_q": (_u8(D, K), _f32(D)), "encoder": (_u8(2 * D, K_ENC), _f32(2 * D)), + "block0_o": (_u8(K, D), _f32(K)), "block0_ffn0": (_u8(FF, K), _f32(FF)), "block0_ffn2": (_u8(K, FF), _f32(K)), + "block1_qkv": (_u8(3 * D, K), _f32(3 * D)), + "block1_o": (_u8(K, D), _f32(K)), "block1_ffn0": (_u8(FF, K), _f32(FF)), "block1_ffn2": (_u8(K, FF), _f32(K)), + } + self.adaln = [(_u8(2 * K, K), _f32(2 * K, 1e-3, 1e-2)) for _ in range(2)] + self.dit_sq = {"block0_q": _f32(K), "encoder": _f32(K_ENC), "block0_o": _f32(D), "block0_ffn0": _f32(K), "block0_ffn2": _f32(FF), + "block1_qkv": _f32(K), "block1_o": _f32(D), "block1_ffn0": _f32(K), "block1_ffn2": _f32(FF)} + + def folded_gamma(self, i, key): + norm = "input_layernorm" if key == "qkv" else "post_attention_layernorm" + return (self.gains[f"L{i}_{norm}"].float() / torch.from_numpy(self.sq[f"L{i}_{key}"])).to(torch.bfloat16) + + # -- the base checkpoint --------------------------------------------------------- + def write_checkpoint(self, root: pathlib.Path): + t = {} + for i in range(2): + t[f"{LM_ROOT}.layers.{i}.self_attn.k_proj.weight"] = torch.zeros(KV, K, dtype=torch.bfloat16) + t[f"{LM_ROOT}.layers.{i}.input_layernorm.weight"] = self.gains[f"L{i}_input_layernorm"] + t[f"{LM_ROOT}.layers.{i}.post_attention_layernorm.weight"] = self.gains[f"L{i}_post_attention_layernorm"] + t[f"{AHK}.model.transformer_blocks.{i}.attn1.to_k.weight"] = torch.zeros(D, K if i else K_ENC, dtype=torch.bfloat16) + root.mkdir(parents=True, exist_ok=True) + save_pt(t, str(root / "model.safetensors")) + return root + + # -- FoldQuantVLA quantized checkpoint (foldquant.quantized_checkpoint) -------------- + LLM_PROJ = {"qkv": [("self_attn.q_proj", K), ("self_attn.k_proj", KV), ("self_attn.v_proj", KV)], + "o": [("self_attn.o_proj", K)], "gateup": [("mlp.gate_proj", INNER), ("mlp.up_proj", INNER)], + "down": [("mlp.down_proj", K)]} + + def write_quantized_checkpoint(self, root: pathlib.Path, *, gptq_kv: bool = True, extra_sites=None): + """The layout FoldQuantVLA's write_quantized_checkpoint produces: qweight/weight_scale in place of + each quantized .weight, foldquant..sq. vectors, sites naming checkpoint keys by rows.""" + t = {"other.weight": torch.ones(3, dtype=torch.bfloat16)} + sites = {"llm": {}, "dit": {}} + + def put(module, key, bits, groups, codes, scale): + params, r0 = [], 0 + for group in groups: + keys = [] + for head, rows in group: + t[f"{head}.qweight"] = torch.from_numpy(codes[r0:r0 + rows].copy()) + t[f"{head}.weight_scale"] = torch.from_numpy(scale[r0:r0 + rows].copy()) + keys.append(f"{head}.weight") + r0 += rows + params.append(keys) + assert r0 == codes.shape[0] + sites[module][key] = {"bits": bits, "params": params} + + for i in range(2): + L = f"{LM_ROOT}.layers.{i}." + t[L + "input_layernorm.weight"] = self.gains[f"L{i}_input_layernorm"] + t[L + "post_attention_layernorm.weight"] = self.gains[f"L{i}_post_attention_layernorm"] + for site, projs in self.LLM_PROJ.items(): + codes, scale = self.llm[f"L{i}_{site}"] + put("llm", f"llm.rtn.L{i}_{site}", 8, [[(L + p, r) for p, r in projs]], codes, scale) + for key, v in self.sq.items(): + t[f"foldquant.llm.sq.{key}"] = torch.from_numpy(v) + B = f"{AHK}.model.transformer_blocks." + put("dit", "dit.block0_q", 4, [[(B + "0.attn1.to_q", D)]], *self.dit["block0_q"]) + kv = [[(B + "0.attn1.to_k", D), (B + "0.attn1.to_v", D)]] + put("dit", "dit.encoder" if gptq_kv else "dit.rtn.block0_kv", 4, kv, *self.dit["encoder"]) + put("dit", "dit.block1_qkv", 4, [[(B + f"1.attn1.to_{x}", D) for x in "qkv"]], *self.dit["block1_qkv"]) + for i in range(2): + put("dit", f"dit.block{i}_o", 4, [[(B + f"{i}.attn1.to_out.0", K)]], *self.dit[f"block{i}_o"]) + put("dit", f"dit.block{i}_ffn0", 4, [[(B + f"{i}.ff.net.0.proj", FF)]], *self.dit[f"block{i}_ffn0"]) + put("dit", f"dit.block{i}_ffn2", 4, [[(B + f"{i}.ff.net.2", K)]], *self.dit[f"block{i}_ffn2"]) + put("dit", f"dit.rtn.block{i}_adaln", 4, [[(B + f"{i}.norm1.linear", 2 * K)]], *self.adaln[i]) + for key, v in self.dit_sq.items(): + t[f"foldquant.dit.sq.{key}"] = torch.from_numpy(v) + for module, key, bits, groups, codes, scale in (extra_sites or []): + put(module, key, bits, groups, codes, scale) + root.mkdir(parents=True, exist_ok=True) + save_pt(t, str(root / "model.safetensors")) + manifest = {"format": cf.FQC_FORMAT, "format_version": 1, "family": "groot_n1_7", "base": {"model_id": "tiny"}, + "weights": {"files": ["model.safetensors"], "sharded": False}, + "modules": {"llm": {"scheme": "w8a8_sr", "config": {"bits": 8, "act_bits": 8, "rot_bs": 64, "act_clip": 1.0}, + "tensors": [f"sq/{k}" for k in self.sq], "sites": sites["llm"]}, + "dit": {"scheme": "w4a4_shg", "config": {"params": {}}, + "tensors": [f"sq/{k}" for k in self.dit_sq], "sites": sites["dit"]}}} + (root / cf.FQ_MANIFEST).write_text(json.dumps(manifest)) + return root + + # -- FoldQuantVLA state ----------------------------------------------------------- + def write_state(self, root: pathlib.Path, llm_scheme="w8a8_sr", dit_scheme="w4a4_shg"): + t = {} + for key, (codes, scale) in self.llm.items(): + t[f"llm/g/llm.rtn.{key}/0/codes"], t[f"llm/g/llm.rtn.{key}/0/scale"] = codes, scale + for key, v in self.sq.items(): + t[f"llm/t/sq/{key}"] = v + for key, (codes, scale) in self.dit.items(): + group = "dit.encoder" if key == "encoder" else f"dit.{key}" + t[f"dit/g/{group}/0/codes"], t[f"dit/g/{group}/0/scale"] = codes, scale + for key, v in self.dit_sq.items(): + t[f"dit/t/sq/{key}"] = v + t["dit/g/dit.rtn/0/codes"], t["dit/g/dit.rtn/0/scale"] = _u8(2 * K, K), _f32(2 * K) # adaLN, by call order + self.write_checkpoint(root) + save_file(t, str(root / cf.FQ_TENSORS)) + manifest = {"format": cf.FQ_FORMAT, "format_version": 2, "family": "groot_n1_7", "base": {"model_id": "tiny"}, + "modules": {"llm": {"scheme": llm_scheme, "config": {"bits": 8, "act_bits": 8, "rot_bs": 64, "act_clip": 1.0}}, + "dit": {"scheme": dit_scheme, "config": {"params": {}}}}} + (root / cf.FQ_MANIFEST).write_text(json.dumps(manifest)) + return root + + def write_state_rtn(self, root: pathlib.Path): + """The same DiT recorded the way the INT8 emitter records it: `dit.rtn/{n}` in call order.""" + t = {} + for key, (codes, scale) in self.llm.items(): + t[f"llm/g/llm.rtn.{key}/0/codes"], t[f"llm/g/llm.rtn.{key}/0/scale"] = codes, scale + for key, v in self.sq.items(): + t[f"llm/t/sq/{key}"] = v + order = [("block0_q",), ("block0_o",), ("encoder",), ("block0_ffn0",), ("block0_ffn2",), + ("block1_qkv",), ("block1_o",), ("block1_ffn0",), ("block1_ffn2",)] + for n, (key,) in enumerate(order): + codes, scale = self.dit[key] + t[f"dit/g/dit.rtn/{n}/codes"], t[f"dit/g/dit.rtn/{n}/scale"] = codes, scale + for key, v in self.dit_sq.items(): + t[f"dit/t/sq/{key}"] = v + self.write_checkpoint(root) + save_file(t, str(root / cf.FQ_TENSORS)) + manifest = {"format": cf.FQ_FORMAT, "format_version": 2, "family": "groot_n1_7", "base": {"model_id": "tiny"}, + "modules": {"llm": {"scheme": "w8a8_sr", "config": {"bits": 8, "act_bits": 8, "rot_bs": 64, "act_clip": 1.0}}, + "dit": {"scheme": "w8a8_sh", "config": {"params": {}}}}} + (root / cf.FQ_MANIFEST).write_text(json.dumps(manifest)) + return root + + # -- plugin ONNX graphs ------------------------------------------------------------ + def write_graphs(self, root: pathlib.Path): + import onnx + from onnx import helper as oh + + root.mkdir(parents=True, exist_ok=True) + + def graph(nodes, name): + g = oh.make_graph(nodes, name, [oh.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1])], + [oh.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1])]) + onnx.save(oh.make_model(g, opset_imports=[oh.make_opsetid("", 17), oh.make_opsetid("trt.plugins", 1)]), str(root / f"{name}.onnx")) + + llm = [] + for i in range(2): + for key, op in (("qkv", "FusedRmsNormLinearInt8"), ("o", "PerRowInt8LinearResidual"), + ("gateup", "FusedRmsNormLinearInt8"), ("down", "PerRowInt8LinearResidual")): + codes, scale = self.llm[f"L{i}_{key}"] + attrs = dict(K=codes.shape[1], N=codes.shape[0], rot_block_size=64, weight_i8=codes.tobytes(), weight_scale=scale.tobytes()) + if key in ("qkv", "gateup"): + attrs["gamma"] = self.folded_gamma(i, key).view(torch.int16).numpy().tobytes() + llm.append(oh.make_node(op, ["x"], ["y"], name=cf.LLM_NODES[key].format(i=i), domain="trt.plugins", **attrs)) + graph(llm, "llm") + + def w(key, stem): + codes, scale = self.dit[key] + return {f"{stem}_i4": codes.tobytes(), f"{stem}_scale": scale.tobytes()} + + dit = [oh.make_node("EncoderPreQuantInt4", ["x"], ["y"], name="encoder_prequant_int4", domain="trt.plugins", + K_enc=K_ENC, act_scale_pre_enc=self.dit_sq["encoder"].tobytes(), block_size=64)] + dit.append(oh.make_node("FusedCrossAttnFullInt4", ["x"], ["y"], name="block0_crossattn_int4", domain="trt.plugins", + K=K, K_enc=K_ENC, inner_dim=D, act_scale_pre_in=self.dit_sq["block0_q"].tobytes(), + act_scale_pre_o=self.dit_sq["block0_o"].tobytes(), **w("block0_q", "weight_q"), + **w("encoder", "weight_kv"), **w("block0_o", "weight_o"))) + dit.append(oh.make_node("FusedSelfAttnFullInt4", ["x"], ["y"], name="block1_selfattn_int4", domain="trt.plugins", + K=K, inner_dim=D, act_scale_pre_in=self.dit_sq["block1_qkv"].tobytes(), + act_scale_pre_o=self.dit_sq["block1_o"].tobytes(), **w("block1_qkv", "weight_qkv"), **w("block1_o", "weight_o"))) + for i in range(2): + dit.append(oh.make_node("FusedFfnBlockInt4", ["x"], ["y"], name=f"block{i}_ffn_int4", domain="trt.plugins", + K=K, inner_dim=FF, block_size=64, act_scale_pre0=self.dit_sq[f"block{i}_ffn0"].tobytes(), + act_scale_pre2=self.dit_sq[f"block{i}_ffn2"].tobytes(), + **w(f"block{i}_ffn0", "weight_proj0"), **w(f"block{i}_ffn2", "weight_proj2"))) + graph(dit, "dit") + return root + + +CONV = types.SimpleNamespace(LM_ROOT=LM_ROOT, AHK=AHK) +FAM = cf.FAMILIES["groot_n1_7"] + + +@pytest.fixture(scope="module") +def tiny(tmp_path_factory): + base = tmp_path_factory.mktemp("tiny") + t = Tiny() + t.state = t.write_state(base / "state") + t.ckpt = t.write_checkpoint(base / "ckpt") + t.graphs = t.write_graphs(base / "onnx") + t.fqc = t.write_quantized_checkpoint(base / "fqc") + t.fqc_rtn = t.write_quantized_checkpoint(base / "fqc_rtn", gptq_kv=False) + return t + + +def test_fold_gain_matches_the_emitter_rule() -> None: + torch.manual_seed(0) + w = (torch.randn(128) * 0.1).to(torch.bfloat16) + s = np.random.default_rng(0).uniform(0.5, 2.0, 128).astype(np.float32) + got = cf._fold_gain(w, s, gemma=True) + want = ((w.float() + 1.0) / torch.from_numpy(s)).to(torch.bfloat16) # the 1 in fp32, one rounding + assert got.dtype == torch.bfloat16 and torch.equal(got, want) + twice = ((w + 1.0).float() / torch.from_numpy(s)).to(torch.bfloat16) # the 1 in bf16: two roundings + assert not torch.equal(got, twice) + f32 = cf._fold_gain(w.float(), s, gemma=False) + assert f32.dtype == torch.float32 and torch.equal(f32, w.float() / torch.from_numpy(s)) + + +def test_state_sites_map_by_name(tiny) -> None: + src = cf.open_source(tiny.state) + assert isinstance(src, cf.FoldQuantState) and src.family == "groot_n1_7" + sites, overrides, kv, extra = cf.build_sites(src, FAM, cf.Checkpoint(tiny.ckpt), CONV) + # 2 layers x 7 LLM projections + block0 (q,k,v,o,ff0,ff2) + block1 (q,k,v,o,ff0,ff2) + assert len(sites) == 14 + 12 and len(overrides) == 4 + q, k, v = (sites[f"vlm.blk.0.attn_{x}.weight"] for x in "qkv") + codes, scale = tiny.llm["L0_qkv"] + assert np.array_equal(np.concatenate([q.codes, k.codes, v.codes]), codes) and np.array_equal(v.wscale, scale[-KV:]) + assert q.bits == 8 and q.ascale is None + # cross block: kv from the encoder record, with the shared encoder SmoothQuant vector + kk = sites["aex.dit.0.attn_k.weight"] + assert kk.bits == 4 and np.array_equal(kk.codes, tiny.dit["encoder"][0][:D]) and np.array_equal(kk.ascale, tiny.dit_sq["encoder"]) + assert np.array_equal(sites["aex.dit.0.attn_q.weight"].ascale, tiny.dit_sq["block0_q"]) + # self block: q/k/v split from the merged record, sharing its vector + vv = sites["aex.dit.1.attn_v.weight"] + assert np.array_equal(vv.codes, tiny.dit["block1_qkv"][0][2 * D:]) and np.array_equal(vv.ascale, tiny.dit_sq["block1_qkv"]) + # folded gains: w / s in fp32, one bf16 rounding; GR00T files keep them as they are (no Gemma - 1) + g = torch.from_numpy(overrides["vlm.blk.1.ffn_norm.weight"]) + assert torch.equal(g.to(torch.bfloat16), tiny.folded_gamma(1, "gateup")) + assert kv["scheme_llm"] == ("str", "w8a8_sr") and kv["action_weight_bits"] == ("u32", 4) and kv["llm_weight_bits"] == ("u32", 8) + assert kv["site_bits"] == ("str", "") and kv["applied_at"] == ("str", "foldquant") + + +def test_graphs_and_state_agree(tiny) -> None: + src = cf.open_source(tiny.state) + ckpt = cf.Checkpoint(tiny.ckpt) + sites, _, _, extra = cf.build_sites(src, FAM, ckpt, CONV) + n = cf.check_against_graphs(src, FAM, ckpt, CONV, sites, extra["gains"], tiny.graphs) + assert n == len(sites) + # a single flipped code is caught + sites["aex.dit.1.ff2.weight"].codes[0, 0] ^= 0xFF + with pytest.raises(SystemExit, match="ff2.weight codes"): + cf.check_against_graphs(src, FAM, ckpt, CONV, sites, extra["gains"], tiny.graphs) + + +def test_call_order_dit_packs_map_by_emitter_order(tiny, tmp_path) -> None: + t = Tiny() + t.__dict__.update({k: v for k, v in tiny.__dict__.items() if k in ("llm", "sq", "gains", "dit", "dit_sq")}) + src = cf.open_source(t.write_state_rtn(tmp_path / "rtn")) + ckpt = cf.Checkpoint(tiny.ckpt) + sites, _, kv, extra = cf.build_sites(src, FAM, ckpt, CONV) + ref, _, _, _ = cf.build_sites(cf.open_source(tiny.state), FAM, ckpt, CONV) + assert set(sites) == set(ref) + for name in sites: + assert np.array_equal(sites[name].codes, ref[name].codes) and np.array_equal(sites[name].ascale, ref[name].ascale) + assert kv["scheme_action"] == ("str", "w8a8_sh") + n = cf.check_against_graphs(src, FAM, ckpt, CONV, sites, extra["gains"], tiny.graphs) + assert n == len(sites) + + +def test_refusals(tiny, tmp_path) -> None: + t = Tiny() + with pytest.raises(SystemExit, match="dense learned rotation"): + cf._check_schemes(cf.open_source(t.write_state(tmp_path / "sr", dit_scheme="w4a4_sr")), "dit") + with pytest.raises(SystemExit, match="no vla.cpp kernel"): + cf._check_schemes(cf.open_source(t.write_state(tmp_path / "plain", llm_scheme="w8a8")), "llm") + with pytest.raises(SystemExit, match="not a FoldQuantVLA quantized model"): + cf.open_source(tmp_path) + src = cf.open_source(tiny.state) + with pytest.raises(SystemExit, match="call order only"): + src.site("dit", "block0_adaln") + + +@pytest.mark.parametrize("layout", ["fqc", "fqc_rtn"]) +def test_quantized_checkpoint_sites_match_the_state(tiny, layout) -> None: + root = getattr(tiny, layout) + src = cf.open_source(root) + assert isinstance(src, cf.FoldQuantCheckpoint) and src.family == "groot_n1_7" + ckpt = cf.Checkpoint(root) + sites, overrides, kv, extra = cf.build_sites(src, FAM, ckpt, CONV) + ref, ref_over, ref_kv, _ = cf.build_sites(cf.open_source(tiny.state), FAM, cf.Checkpoint(tiny.ckpt), CONV) + assert set(sites) == set(ref) + for name in sites: + a, b = sites[name], ref[name] + assert a.bits == b.bits and np.array_equal(a.codes, b.codes) and np.array_equal(a.wscale, b.wscale), name + assert (a.ascale is None and b.ascale is None) or np.array_equal(a.ascale, b.ascale), name + for name in overrides: + assert np.array_equal(overrides[name], ref_over[name]), name + assert kv["scheme_llm"] == ref_kv["scheme_llm"] and kv["action_weight_bits"] == ("u32", 4) + assert "adaLN dequantized from INT4" in kv["provenance"][1] + src.check_accounting() # every projection is a site or the allowed float adaLN + assert cf.check_against_graphs(src, FAM, ckpt, CONV, sites, extra["gains"], tiny.graphs) == len(sites) + + +def test_quantized_checkpoint_base_view(tiny) -> None: + ckpt = cf.Checkpoint(tiny.fqc) + keys = set(ckpt.keys()) + B = f"{AHK}.model.transformer_blocks." + assert ckpt.is_quantized and "other.weight" in keys + assert not any(k.endswith((".qweight", ".weight_scale")) or k.startswith("foldquant.") for k in keys) + assert f"{LM_ROOT}.layers.0.self_attn.k_proj.weight" in keys and f"{B}0.norm1.linear.weight" in keys + assert ckpt.shape(f"{LM_ROOT}.layers.0.self_attn.k_proj.weight") == (KV, K) + assert ckpt.shape(f"{B}0.attn1.to_k.weight") == (D, K_ENC) # int4: K from the packed width + # adaLN: INT4 codes x the scale rounded to bf16 (what AdaLNModInt4 multiplies by), one bf16 rounding + codes, scale = tiny.adaln[0] + lo = (codes & 0xF).astype(np.int8) + hi = (codes >> 4).astype(np.int8) + q = np.stack([np.where(lo >= 8, lo - 16, lo), np.where(hi >= 8, hi - 16, hi)], -1).reshape(codes.shape[0], -1) + s_bf16 = torch.from_numpy(scale).to(torch.bfloat16).float().numpy() + want = torch.from_numpy(q.astype(np.float32) * s_bf16[:, None]).to(torch.bfloat16) + got = ckpt.tensor(f"{B}0.norm1.linear.weight") + assert got.dtype == torch.bfloat16 and torch.equal(got, want) + m = ckpt.mapping((f"{AHK}.",)) + assert f"{B}1.norm1.linear.weight" in m and "other.weight" not in m and len(m) == len(list(m)) + h = ckpt.handle() + assert h.get_slice(f"{LM_ROOT}.layers.0.mlp.gate_proj.weight").get_shape() == [INNER, K] + + +def test_base_view_points_the_converter_at_the_quantized_checkpoint(tiny, tmp_path) -> None: + import safetensors + + ckpt = cf.Checkpoint(tiny.fqc) + other = tmp_path / "other.safetensors" + save_pt({"x": torch.zeros(2)}, str(other)) + conv = types.SimpleNamespace(load_safetensors=lambda path, keep=None: "real", safe_open=safetensors.safe_open) + real_load, real_open = conv.load_safetensors, conv.safe_open + with cf.base_view(conv, ckpt): + assert isinstance(conv.load_safetensors(tiny.fqc), cf._TensorMap) + assert conv.load_safetensors(tmp_path) == "real" + assert isinstance(conv.safe_open(tiny.fqc / "model.safetensors", framework="pt"), cf._Handle) + with conv.safe_open(str(other), framework="pt") as f: + assert list(f.keys()) == ["x"] + assert conv.load_safetensors is real_load and conv.safe_open is real_open + plain = cf.Checkpoint(tiny.ckpt) # an ordinary checkpoint: nothing is patched + with cf.base_view(conv, plain): + assert conv.load_safetensors is real_load + + +def test_quantized_checkpoint_refusals(tiny, tmp_path) -> None: + t = Tiny() + t.__dict__.update({k: v for k, v in tiny.__dict__.items() if k in ("llm", "sq", "gains", "dit", "dit_sq", "adaln")}) + # a site vla.cpp has no mapping for + root = t.write_quantized_checkpoint(tmp_path / "bad_site", extra_sites=[ + ("dit", "dit.block0_mystery", 4, [[(f"{AHK}.x", 64)]], _u8(64, 128), _f32(64))]) + with pytest.raises(SystemExit, match="no mapping for"): + cf.open_source(root) + # a quantized projection no GGUF site consumes (layer 7 has an `o` site but no layers otherwise) + root = t.write_quantized_checkpoint(tmp_path / "orphan", extra_sites=[ + ("llm", "llm.rtn.L7_o", 8, [[(f"{LM_ROOT}.layers.7.self_attn.o_proj", K)]], _i8(K, K), _f32(K))]) + src = cf.open_source(root) + cf.build_sites(src, FAM, cf.Checkpoint(root), CONV) + with pytest.raises(SystemExit, match="neither a vla.cpp site"): + src.check_accounting() + # a newer format version, and an unknown FoldQuant format + m = json.loads((tiny.fqc / cf.FQ_MANIFEST).read_text()) + for fmt, ver, match in ((cf.FQC_FORMAT, 2, "newer than this converter"), ("foldquant-something", 1, "not one this converter reads")): + d = tmp_path / f"m_{ver}_{fmt}" + d.mkdir() + (d / "model.safetensors").symlink_to(tiny.fqc / "model.safetensors") + (d / cf.FQ_MANIFEST).write_text(json.dumps({**m, "format": fmt, "format_version": ver})) + with pytest.raises(SystemExit, match=match): + cf.open_source(d) + + +def test_writer_round_trip(tmp_path: pathlib.Path) -> None: + out = tmp_path / "t.gguf" + codes4 = np.random.default_rng(0).integers(0, 255, (64, 64), dtype=np.uint8) # K = 128, nibble-packed + sites = {"blk.0.w.weight": Site(codes=codes4, wscale=np.full(64, 0.5, np.float32), + ascale=np.full(128, 2.0, np.float32), bits=4)} + overrides = {"blk.0.norm.weight": np.full(128, 0.25, np.float32)} + w = QuantizingGGUFWriter(out, "t", sites=sites, overrides=overrides, + quant_kv={"method": ("str", "foldquant"), "llm_weight_bits": ("u32", 4)}, f32_overrides=True) + w.add_tensor("blk.0.w.weight", np.zeros((64, 128), np.float32)) + w.add_tensor("blk.0.norm.weight", np.zeros(128, np.float32)) + w.add_tensor("other", np.ones(3, np.float32)) + assert w.unconsumed() == [] + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + r = gguf.GGUFReader(str(out)) + t = {x.name: x for x in r.tensors} + assert t["blk.0.w.weight"].tensor_type == gguf.GGMLQuantizationType.I8 + assert np.array_equal(np.asarray(t["blk.0.w.weight"].data).view(np.uint8).reshape(64, 64), codes4) + assert np.allclose(t["blk.0.w.wscale"].data, 0.5) and np.allclose(t["blk.0.w.ascale"].data, 2.0) + assert t["blk.0.norm.weight"].tensor_type == gguf.GGMLQuantizationType.F32 + assert np.allclose(t["blk.0.norm.weight"].data, 0.25) + assert "t.quant.method" in r.fields + + +def test_writer_rejects_shape_mismatch(tmp_path: pathlib.Path) -> None: + sites = {"a.weight": Site(codes=np.zeros((64, 64), np.int8), wscale=np.ones(64, np.float32), ascale=None, bits=8)} + w = QuantizingGGUFWriter(tmp_path / "t.gguf", "t", sites=sites, overrides={}, quant_kv={}) + with pytest.raises(ValueError, match="the site is"): + w.add_tensor("a.weight", np.zeros((64, 128), np.float32)) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-q", "-p", "no:cacheprovider"])) diff --git a/tests/test_foldquant_cpu_op.cpp b/tests/test_foldquant_cpu_op.cpp new file mode 100644 index 0000000..213f0f0 --- /dev/null +++ b/tests/test_foldquant_cpu_op.cpp @@ -0,0 +1,294 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant CPU reference ops, run as GGML_OP_CUSTOM nodes on the CPU backend +// with several threads, checked three ways: +// 1. bit-identical to the reference functions called row by row (graph +// plumbing, src slots, thread partitioning); +// 2. within +-1 code of an independent double-precision dense-Hadamard +// implementation (the math itself); +// 3. a golden FNV-1a checksum of the codes for the fixed LCG input, shared +// with tests/py/test_foldquant_ref.py so the numpy reference and the C++ +// one cannot drift apart. + +#include "foldquant.h" +#include "foldquant_ref.h" +#include "layers/fq_linear.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include + +namespace { + +struct Lcg { + uint32_t s; + explicit Lcg(uint32_t seed) : s(seed) {} + // Same generator as scripts/foldquant_ref.py: value in [-1, 1]. + float next() { + s = s * 1664525u + 1013904223u; + return ((float) ((s >> 8) & 0xFFFFu) / 65535.0f) * 2.0f - 1.0f; + } +}; + +uint64_t fnv1a(const uint8_t * p, size_t n) { + uint64_t h = 14695981039346656037ull; + for (size_t i = 0; i < n; ++i) { h ^= p[i]; h *= 1099511628211ull; } + return h; +} + +struct Case { + const char * name; + int wbits, abits, rot; + bool gamma, ascale, fold_before, bias; +}; + +constexpr int64_t K = 128, N = 64, T = 5; + +// FNV-1a of the INT8 codes of the first case (no norm, rot 64) for the LCG input. +constexpr uint64_t FQ_GOLDEN_FNV = 0xb97c1c173d2d0278ull; + +// Independent double-precision reference of the activation transform. +void ref_act_double(const std::vector & x, const std::vector & as, const std::vector & ga, + const vla::FqActSpec & s, std::vector & codes, std::vector & scales) { + codes.assign((size_t) T * K, 0); + scales.assign((size_t) T, 0.0); + const int bs = s.rot_block; + std::vector H((size_t) bs * bs, 0.0); + for (int i = 0; i < bs; ++i) + for (int j = 0; j < bs; ++j) { + int bits = i & j, par = 0; + while (bits) { par ^= bits & 1; bits >>= 1; } + H[(size_t) i * bs + j] = (par ? -1.0 : 1.0) / std::sqrt((double) bs); + } + const double qmax = s.abits == 4 ? 7.0 : 127.0; + for (int64_t t = 0; t < T; ++t) { + std::vector y(K); + for (int64_t k = 0; k < K; ++k) y[k] = x[(size_t) t * K + k]; + if (s.has_gamma) { + double ss = 0.0; + for (int64_t k = 0; k < K; ++k) ss += y[k] * y[k]; + const double rstd = 1.0 / std::sqrt(ss / (double) K + (double) s.eps); + for (int64_t k = 0; k < K; ++k) y[k] = y[k] * rstd * ga[k]; + } + if (!as.empty() && s.fold_before) for (int64_t k = 0; k < K; ++k) y[k] /= as[k]; + if (bs > 1) { + std::vector z(K); + for (int64_t b = 0; b < K; b += bs) + for (int i = 0; i < bs; ++i) { + double acc = 0.0; + for (int j = 0; j < bs; ++j) acc += H[(size_t) i * bs + j] * y[b + j]; + z[b + i] = acc; + } + y = z; + } + if (!as.empty() && !s.fold_before) for (int64_t k = 0; k < K; ++k) y[k] /= as[k]; + double amax = 0.0; + for (int64_t k = 0; k < K; ++k) amax = std::fmax(amax, std::fabs(y[k])); + double scale = (double) s.clip * amax / qmax; + if (scale < 1e-12) scale = 1e-12; + scales[t] = scale; + for (int64_t k = 0; k < K; ++k) { + double q = std::nearbyint(y[k] / scale); + q = std::fmin(qmax, std::fmax(-qmax, q)); + codes[(size_t) t * K + k] = (int) q; + } + } +} + +int run_case(ggml_backend_t backend, const Case & c, bool print_hash) { + std::printf("case %-22s W%dA%d rot%-3d gamma=%d ascale=%d before=%d bias=%d\n", + c.name, c.wbits, c.abits, c.rot, c.gamma, c.ascale, c.fold_before, c.bias); + + Lcg rng(0x5eed1234u); + std::vector hx((size_t) K * T), hw((size_t) K * N), hws(N), hb(N), has(K), hga(K); + for (auto & v : hx) v = rng.next() * 4.0f; + for (auto & v : hw) v = rng.next(); + for (auto & v : hws) v = 0.01f + 0.02f * std::fabs(rng.next()); + for (auto & v : hb) v = rng.next() * 0.5f; + for (auto & v : has) v = 0.5f + std::fabs(rng.next()); + for (auto & v : hga) v = 0.75f + 0.5f * std::fabs(rng.next()); + + // Weight codes: per-output-row symmetric of the float w (rounded here; the + // GEMM only needs *some* valid codes). + const float wq = c.wbits == 4 ? 7.0f : 127.0f; + std::vector wcodes((size_t) K * N); + for (int64_t n = 0; n < N; ++n) { + float amax = 0.f; + for (int64_t k = 0; k < K; ++k) amax = std::fmax(amax, std::fabs(hw[(size_t) n * K + k])); + const float sc = amax / wq; + for (int64_t k = 0; k < K; ++k) + wcodes[(size_t) n * K + k] = (int8_t) std::nearbyint(hw[(size_t) n * K + k] / sc); + } + const int64_t kpw = vla::fq_w_kpack(K, c.wbits); + std::vector wpacked((size_t) kpw * N); + if (c.wbits == 4) { + for (int64_t n = 0; n < N; ++n) + for (int64_t k = 0; k < K; k += 2) + wpacked[(size_t) n * kpw + k / 2] = (int8_t) ((wcodes[(size_t) n * K + k] & 0xF) | ((wcodes[(size_t) n * K + k + 1] & 0xF) << 4)); + } else { + wpacked = wcodes; + } + + ggml_init_params ip = { (size_t) 64 * 1024 * 1024, nullptr, false }; + ggml_context * C = ggml_init(ip); + + ggml_tensor * x = ggml_new_tensor_2d(C, GGML_TYPE_F32, K, T); + ggml_tensor * w = ggml_new_tensor_2d(C, GGML_TYPE_I8, kpw, N); + ggml_tensor * ws = ggml_new_tensor_1d(C, GGML_TYPE_F32, N); + ggml_tensor * b = c.bias ? ggml_new_tensor_1d(C, GGML_TYPE_F32, N) : nullptr; + ggml_tensor * as = c.ascale ? ggml_new_tensor_1d(C, GGML_TYPE_F32, K) : nullptr; + ggml_tensor * ga = c.gamma ? ggml_new_tensor_1d(C, GGML_TYPE_F32, K) : nullptr; + ggml_set_name(w, "site"); + std::memcpy(x->data, hx.data(), ggml_nbytes(x)); + std::memcpy(w->data, wpacked.data(), ggml_nbytes(w)); + std::memcpy(ws->data, hws.data(), ggml_nbytes(ws)); + if (b) std::memcpy(b->data, hb.data(), ggml_nbytes(b)); + if (as) std::memcpy(as->data, has.data(), ggml_nbytes(as)); + if (ga) std::memcpy(ga->data, hga.data(), ggml_nbytes(ga)); + + vla::FqLinear s; + s.w = w; s.wscale = ws; s.bias = b; s.ascale = as; s.gamma = ga; + s.act.K = K; s.act.abits = c.abits; s.act.rot_block = c.rot; s.act.fold_before = c.fold_before; + s.act.has_gamma = c.gamma; s.act.has_ascale = c.ascale; s.act.clip = c.abits == 4 ? 0.9f : 1.0f; s.act.eps = 1e-6f; + s.gemm.K = K; s.gemm.N = N; s.gemm.wbits = c.wbits; + + ggml_tensor * xq = vla::fq_act(C, s, x); + ggml_tensor * y = vla::fq_gemm(C, s, xq); + ggml_cgraph * gf = ggml_new_graph(C); + ggml_build_forward_expand(gf, y); + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + std::printf("FAIL: graph compute\n"); + return 1; + } + + // 1. bit-identical to the reference functions. + const int64_t rb = vla::fq_act_row_bytes(K, c.abits); + const int64_t kp = vla::fq_act_kpack(K, c.abits); + std::vector blob((size_t) rb * T); + { + std::vector tmp(K); float partial[vla::fqref::NT]; + for (int64_t t = 0; t < T; ++t) { + uint8_t * row = blob.data() + (size_t) t * rb; + std::memset(row, 0, rb); + vla::fqref::act_row(hx.data() + (size_t) t * K, as ? has.data() : nullptr, ga ? hga.data() : nullptr, + s.act, tmp.data(), partial, (int8_t *) row, (float *) (row + kp)); + } + } + for (int64_t t = 0; t < T; ++t) { + const uint8_t * got = (const uint8_t *) xq->data + (size_t) t * rb; + if (std::memcmp(got, blob.data() + (size_t) t * rb, kp + 4) != 0) { + std::printf("FAIL: row %lld of the act blob differs from the reference\n", (long long) t); + return 1; + } + } + std::vector xcodes((size_t) T * K); + for (int64_t t = 0; t < T; ++t) { + const int8_t * row = (const int8_t *) xq->data + (size_t) t * rb; + if (c.abits == 4) vla::fqref::unpack_nibbles(row, K, xcodes.data() + (size_t) t * K); + else std::memcpy(xcodes.data() + (size_t) t * K, row, K); + } + for (int64_t t = 0; t < T; ++t) { + float xs; std::memcpy(&xs, (const uint8_t *) xq->data + (size_t) t * rb + kp, 4); + for (int64_t n = 0; n < N; ++n) { + const float want = vla::fqref::gemm_dot(wcodes.data() + (size_t) n * K, xcodes.data() + (size_t) t * K, + K, xs, hws[n], b ? hb[n] : 0.0f); + const float got = ((const float *) y->data)[(size_t) t * N + n]; + if (std::memcmp(&want, &got, 4) != 0) { + std::printf("FAIL: y[%lld][%lld] = %g, reference %g\n", (long long) t, (long long) n, got, want); + return 1; + } + } + } + + // 2. the math, against double precision. + std::vector dcodes; std::vector dscales; + ref_act_double(hx, as ? has : std::vector(), hga, s.act, dcodes, dscales); + int off_by_one = 0; + for (int64_t t = 0; t < T; ++t) { + float xs; std::memcpy(&xs, (const uint8_t *) xq->data + (size_t) t * rb + kp, 4); + if (std::fabs((double) xs - dscales[t]) > 1e-5 * dscales[t]) { + std::printf("FAIL: row %lld scale %g vs double %g\n", (long long) t, xs, dscales[t]); + return 1; + } + for (int64_t k = 0; k < K; ++k) { + const int d = std::abs((int) xcodes[(size_t) t * K + k] - dcodes[(size_t) t * K + k]); + if (d > 1) { std::printf("FAIL: code [%lld][%lld] off by %d\n", (long long) t, (long long) k, d); return 1; } + off_by_one += d; + } + } + if (off_by_one > (int) (T * K) / 100) { + std::printf("FAIL: %d codes differ from the double reference by one\n", off_by_one); + return 1; + } + + if (print_hash) + std::printf(" codes fnv1a = 0x%016llx\n", (unsigned long long) fnv1a((const uint8_t *) xcodes.data(), xcodes.size())); + std::printf(" ok (%d codes off by one vs double)\n", off_by_one); + ggml_free(C); + return 0; +} + +} // namespace + +int main() { + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_cpu_set_n_threads(backend, 4); + + const Case cases[] = { + { "w8a8_rot64", 8, 8, 64, false, false, false, true }, + { "w8a8_rot64_gamma", 8, 8, 64, true, false, false, false }, + { "w8a8_rot32_pre", 8, 8, 32, false, true, true, true }, + { "w8a8_rot64_post", 8, 8, 64, false, true, false, false }, + { "w8a8_norot", 8, 8, 1, false, false, false, false }, + { "w4a8_rot64_gamma", 4, 8, 64, true, false, false, true }, + { "w4a4_rot64_pre", 4, 4, 64, false, true, true, true }, + }; + int i = 0; + for (const Case & c : cases) { + if (run_case(backend, c, i == 0)) return 1; + ++i; + } + + // 3. golden checksum of the first case (no gamma: its reduction order is the + // one thing numpy cannot mirror exactly). + { + Lcg rng(0x5eed1234u); + std::vector hx((size_t) K * T); + for (auto & v : hx) v = rng.next() * 4.0f; + vla::FqActSpec a; a.K = K; a.abits = 8; a.rot_block = 64; + std::vector codes((size_t) T * K); + std::vector tmp(K); float partial[vla::fqref::NT]; float sc; + for (int64_t t = 0; t < T; ++t) + vla::fqref::act_row(hx.data() + (size_t) t * K, nullptr, nullptr, a, tmp.data(), partial, + codes.data() + (size_t) t * K, &sc); + const uint64_t h = fnv1a((const uint8_t *) codes.data(), codes.size()); + const uint64_t golden = FQ_GOLDEN_FNV; // pinned by tests/py/test_foldquant_ref.py too + if (h != golden) { + std::printf("FAIL: golden checksum 0x%016llx, expected 0x%016llx\n", + (unsigned long long) h, (unsigned long long) golden); + return 1; + } + } + + std::printf("PASS\n"); + ggml_backend_free(backend); + return 0; +} diff --git a/tests/test_foldquant_cuda_op.cpp b/tests/test_foldquant_cuda_op.cpp new file mode 100644 index 0000000..b5be04f --- /dev/null +++ b/tests/test_foldquant_cuda_op.cpp @@ -0,0 +1,279 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant nodes on the CUDA backend versus the CPU reference: the activation +// blob and the F32 GEMM output must be byte-identical (int32 accumulation, +// fixed reduction tree, no FMA contraction on either side). Also covers the +// ggml plumbing the kernels depend on: the custom-op userdata decode, the hook +// dispatcher composing with the BF16 handler, and gallocr sizing of the blob. +// Runs once per ctest registration: plain, and under VLA_FQ_CPU_REF=1 (staging +// shim). Skips itself (exit 0) when no CUDA device is present. + +#include "foldquant.h" +#include "foldquant_ref.h" +#include "layers/fq_linear.h" +#include "act_dtype.h" +#include "cuda/vla_cuda_ops.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml-cuda.h" + +#include +#include +#include +#include +#include + +namespace { + +struct Lcg { + uint32_t s; + explicit Lcg(uint32_t seed) : s(seed) {} + float next() { + s = s * 1664525u + 1013904223u; + return ((float) ((s >> 8) & 0xFFFFu) / 65535.0f) * 2.0f - 1.0f; + } +}; + +struct Shape { int64_t K, N, T; }; + +struct Case { + const char * name; + int wbits, abits, rot; + bool gamma, ascale, fold_before, bias; + bool res = false; // residual fused into the GEMM epilogue + int heads = 0; // >0: head-layout epilogue (hd = 64, 1-3 parts, last part V) +}; + +struct Inputs { + std::vector x, ws, b, as, ga, res; + std::vector w; // packed +}; + +Inputs make_inputs(const Shape & sh, const Case & c, uint32_t seed) { + Inputs in; + Lcg rng(seed); + in.x.resize((size_t) sh.K * sh.T); in.ws.resize(sh.N); in.b.resize(sh.N); in.as.resize(sh.K); in.ga.resize(sh.K); + for (auto & v : in.x) v = rng.next() * 4.0f; + for (auto & v : in.ws) v = 0.01f + 0.02f * std::fabs(rng.next()); + for (auto & v : in.b) v = rng.next() * 0.5f; + in.res.resize((size_t) sh.N * sh.T); + for (auto & v : in.res) v = rng.next() * 3.0f; + for (auto & v : in.as) v = 0.5f + std::fabs(rng.next()); + for (auto & v : in.ga) v = 0.75f + 0.5f * std::fabs(rng.next()); + const int64_t kpw = vla::fq_w_kpack(sh.K, c.wbits); + in.w.resize((size_t) kpw * sh.N); + const int wq = c.wbits == 4 ? 7 : 127; + for (int64_t n = 0; n < sh.N; ++n) { + std::vector codes(sh.K); + for (int64_t k = 0; k < sh.K; ++k) codes[k] = (int) std::nearbyint(rng.next() * (float) wq); + if (c.wbits == 4) + for (int64_t k = 0; k < sh.K; k += 2) + in.w[(size_t) n * kpw + k / 2] = (int8_t) ((codes[k] & 0xF) | ((codes[k + 1] & 0xF) << 4)); + else + for (int64_t k = 0; k < sh.K; ++k) in.w[(size_t) n * kpw + k] = (int8_t) codes[k]; + } + return in; +} + +struct Result { + std::vector blob; + std::vector y; + std::vector bf16_probe; // output of a BF16 mul_mat_t node in the same graph +}; + +Result run(ggml_backend_t backend, const Shape & sh, const Case & c, const Inputs & in, bool with_bf16_probe) { + Result r; + ggml_init_params p = { (size_t) 64 * 1024 * 1024, nullptr, true }; + ggml_context * C = ggml_init(p); + + const int64_t kpw = vla::fq_w_kpack(sh.K, c.wbits); + ggml_tensor * x = ggml_new_tensor_2d(C, GGML_TYPE_F32, sh.K, sh.T); + ggml_tensor * w = ggml_new_tensor_2d(C, GGML_TYPE_I8, kpw, sh.N); + ggml_tensor * ws = ggml_new_tensor_1d(C, GGML_TYPE_F32, sh.N); + ggml_tensor * b = c.bias ? ggml_new_tensor_1d(C, GGML_TYPE_F32, sh.N) : nullptr; + ggml_tensor * rs = c.res ? ggml_new_tensor_2d(C, GGML_TYPE_F32, sh.N, sh.T) : nullptr; + ggml_tensor * as = c.ascale ? ggml_new_tensor_1d(C, GGML_TYPE_F32, sh.K) : nullptr; + ggml_tensor * ga = c.gamma ? ggml_new_tensor_1d(C, GGML_TYPE_F32, sh.K) : nullptr; + ggml_tensor * Wb = with_bf16_probe ? ggml_new_tensor_2d(C, GGML_TYPE_BF16, sh.K, 64) : nullptr; + for (ggml_tensor * t : {x, w, ws, b, as, ga, Wb, rs}) if (t) ggml_set_input(t); + ggml_set_name(w, "site"); + + vla::FqLinear s; + s.w = w; s.wscale = ws; s.bias = b; s.ascale = as; s.gamma = ga; + s.act.K = sh.K; s.act.abits = c.abits; s.act.rot_block = c.rot; s.act.fold_before = c.fold_before; + s.act.has_gamma = c.gamma; s.act.has_ascale = c.ascale; s.act.clip = c.abits == 4 ? 0.9f : 1.0f; s.act.eps = 1e-6f; + s.gemm.K = sh.K; s.gemm.N = sh.N; s.gemm.wbits = c.wbits; + if (c.heads) { // hd = 64; as many parts as the model uses (q/k/v = 3, k/v = 2, one), last part V + const int parts = sh.N % 192 == 0 ? 3 : sh.N % 128 == 0 ? 2 : 1; + if (sh.N % (64 * parts) == 0) + vla::fq_set_heads(s, 64, (int) (sh.N / (64 * parts)), 1u << (parts - 1)); + } + + ggml_tensor * xq = vla::fq_act(C, s, x); + ggml_tensor * y = vla::fq_gemm(C, s, xq, rs); + ggml_set_output(xq); ggml_set_output(y); + ggml_tensor * probe = nullptr; + if (Wb) { + // The BF16 handler must still see its nodes when both are registered. + ggml_tensor * xb = vla::as_type(C, x, GGML_TYPE_BF16); + probe = vla::as_type(C, vla::mm_act(C, Wb, xb, GGML_TYPE_BF16), GGML_TYPE_F32); + ggml_set_output(probe); + } + + ggml_cgraph * gf = ggml_new_graph(C); + ggml_build_forward_expand(gf, y); + if (probe) ggml_build_forward_expand(gf, probe); + + ggml_gallocr_t ga_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ga_ || !ggml_gallocr_alloc_graph(ga_, gf)) { std::fprintf(stderr, "alloc failed\n"); return r; } + + ggml_backend_tensor_set(x, in.x.data(), 0, ggml_nbytes(x)); + ggml_backend_tensor_set(w, in.w.data(), 0, ggml_nbytes(w)); + ggml_backend_tensor_set(ws, in.ws.data(), 0, ggml_nbytes(ws)); + if (b) ggml_backend_tensor_set(b, in.b.data(), 0, ggml_nbytes(b)); + if (rs) ggml_backend_tensor_set(rs, in.res.data(), 0, ggml_nbytes(rs)); + if (as) ggml_backend_tensor_set(as, in.as.data(), 0, ggml_nbytes(as)); + if (ga) ggml_backend_tensor_set(ga, in.ga.data(), 0, ggml_nbytes(ga)); + if (Wb) { + std::vector hw((size_t) sh.K * 64); + Lcg rng(7u); + for (auto & v : hw) v = rng.next(); + std::vector t(hw.size()); + ggml_fp32_to_bf16_row(hw.data(), t.data(), (int64_t) hw.size()); + ggml_backend_tensor_set(Wb, t.data(), 0, ggml_nbytes(Wb)); + } + + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "compute failed\n"); return r; } + + r.blob.resize(ggml_nbytes(xq)); + ggml_backend_tensor_get(xq, r.blob.data(), 0, r.blob.size()); + r.y.resize((size_t) ggml_nelements(y)); + ggml_backend_tensor_get(y, r.y.data(), 0, ggml_nbytes(y)); + if (probe) { + r.bf16_probe.resize((size_t) ggml_nelements(probe)); + ggml_backend_tensor_get(probe, r.bf16_probe.data(), 0, ggml_nbytes(probe)); + } + ggml_gallocr_free(ga_); + ggml_free(C); + return r; +} + +// The CPU op reads its spec through userdata; the CUDA decoder mirrors the +// private ggml_custom_op_params layout. Pin that the pointer round-trips. +bool check_userdata_layout() { + ggml_init_params p = { (size_t) 1024 * 1024, nullptr, true }; + ggml_context * C = ggml_init(p); + ggml_tensor * a = ggml_new_tensor_1d(C, GGML_TYPE_F32, 16); + int marker = 42; + ggml_tensor * args[1] = { a }; + ggml_tensor * t = ggml_custom_4d(C, GGML_TYPE_F32, 16, 1, 1, 1, args, 1, nullptr, GGML_N_TASKS_MAX, &marker); + struct { void * fun; int n_tasks; void * userdata; } mirror; + std::memcpy(&mirror, t->op_params, sizeof(mirror)); + const bool ok = mirror.userdata == &marker && mirror.n_tasks == GGML_N_TASKS_MAX && mirror.fun == nullptr; + ggml_free(C); + return ok; +} + +} // namespace + +int main() { + ggml_backend_t cuda = ggml_backend_cuda_init(0); + if (!cuda) { + std::printf("foldquant_cuda_op: no CUDA device, skipping\n"); + return 0; + } + if (!check_userdata_layout()) { + std::printf("FAIL: ggml_custom_op_params layout changed; update vla_cuda_foldquant.cu\n"); + return 1; + } + vla::cuda_register_bf16_ops(); + vla::cuda_register_foldquant_ops(); + + ggml_backend_t cpu = ggml_backend_cpu_init(); + ggml_backend_cpu_set_n_threads(cpu, 4); + + const Case cases[] = { + { "w8a8_rot64_bias", 8, 8, 64, false, false, false, true }, + { "w8a8_rot64_bias_res", 8, 8, 64, false, false, false, true, true }, + { "w4a4_rot64_res", 4, 4, 64, false, true, true, false, true }, + { "w8a8_rot64_heads", 8, 8, 64, false, false, false, true, false, 4 }, + { "w4a4_rot64_heads", 4, 4, 64, false, true, true, true, false, 2 }, + { "w8a8_rot64_gamma", 8, 8, 64, true, false, false, false }, + { "w8a8_rot32_pre", 8, 8, 32, false, true, true, true }, + { "w8a8_rot64_post", 8, 8, 64, false, true, false, false }, + { "w8a8_norot", 8, 8, 1, false, false, false, false }, + { "w4a8_rot64_gamma", 4, 8, 64, true, false, false, true }, // CPU-reference fallback on CUDA today + { "w4a4_rot64_pre", 4, 4, 64, false, true, true, true }, + }; + const Shape shapes[] = { + { 128, 64, 5 }, + { 2048, 4096, 41 }, // GR00T N1.7 LLM q/k/v at a DiT-sized M + { 2048, 2048, 130 }, // o_proj, M above one tile with a tail + { 6144, 2048, 130 }, // ffn_down + { 1536, 4608, 41 }, // DiT fused qkv + { 2048, 12288, 160 }, // fused gate+up at the N1.7 prefill: the 128x64x128 tile + { 2048, 2048, 193 }, // 192-row tile with a one-row second M tile + { 2048, 4096, 300 }, // prefix past 256 tokens: 128-row tile, three M tiles + // pi0.5 (Gemma-2B prefix, Gemma-300M expert; q and k/v are separate sites) + { 2048, 2048, 700 }, // prefix q / o over a ~700-token prefix + { 2048, 256, 700 }, // prefix k / v (one KV head) + { 2048, 16384, 160 }, // prefix gate / up + { 16384, 2048, 160 }, // prefix down: K past the CTA-per-row prologue + { 1024, 2048, 50 }, // expert q + { 1024, 256, 50 }, // expert k / v + { 2048, 1024, 50 }, // expert o + { 1024, 4096, 50 }, // expert gate / up + { 4096, 1024, 50 }, // expert down + }; + + int n = 0; + for (const Shape & sh : shapes) { + for (const Case & c : cases) { + // Big W4 cases are slow on the host path; the small-M ones (DiT / expert) are cheap. + if (sh.K > 128 && (c.wbits != 8 || c.abits != 8) && sh.T > 64) continue; + const Inputs in = make_inputs(sh, c, 0x1234u + (uint32_t) n); + const Result ref = run(cpu, sh, c, in, false); + const Result got = run(cuda, sh, c, in, true); + if (ref.blob.empty() || got.blob.empty() || ref.y.empty() || got.y.empty()) { + std::printf("FAIL: %s K=%lld N=%lld T=%lld produced no output\n", c.name, (long long) sh.K, (long long) sh.N, (long long) sh.T); + return 1; + } + // Compare codes + scale of every row; the padding bytes are uninitialised. + const int64_t rb = vla::fq_act_row_bytes(sh.K, c.abits), kp = vla::fq_act_kpack(sh.K, c.abits); + for (int64_t t = 0; t < sh.T; ++t) + if (std::memcmp(ref.blob.data() + (size_t) t * rb, got.blob.data() + (size_t) t * rb, kp + 4) != 0) { + std::printf("FAIL: %s K=%lld N=%lld T=%lld: act row %lld differs\n", c.name, + (long long) sh.K, (long long) sh.N, (long long) sh.T, (long long) t); + return 1; + } + if (std::memcmp(ref.y.data(), got.y.data(), ref.y.size() * sizeof(float)) != 0) { + size_t first = 0; + while (first < ref.y.size() && ref.y[first] == got.y[first]) ++first; + std::printf("FAIL: %s K=%lld N=%lld T=%lld: y[%zu] = %.9g vs CPU %.9g\n", c.name, + (long long) sh.K, (long long) sh.N, (long long) sh.T, first, got.y[first], ref.y[first]); + return 1; + } + if (got.bf16_probe.empty()) { std::printf("FAIL: BF16 probe missing\n"); return 1; } + std::printf("ok %-18s K=%-5lld N=%-5lld T=%-4lld\n", c.name, (long long) sh.K, (long long) sh.N, (long long) sh.T); + ++n; + } + } + std::printf("PASS (%d shape/case pairs)\n", n); + ggml_backend_free(cpu); + ggml_backend_free(cuda); + return 0; +} diff --git a/tests/test_foldquant_dequant.cpp b/tests/test_foldquant_dequant.cpp new file mode 100644 index 0000000..fc05d89 --- /dev/null +++ b/tests/test_foldquant_dequant.cpp @@ -0,0 +1,110 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant CPU reference ops, run as GGML_OP_CUSTOM nodes on the CPU backend +// with several threads, checked three ways: +// fq_dequant_rows (FoldQuant read back as a float weight, foldquant.h) against +// the dense product it stands for: W_deq . R . diag(1/a) (fold before) or +// W_deq . diag(1/a) . R (fold after), with R the block-diagonal normalised +// natural-order Sylvester-Hadamard matrix. Both widths, both fold orders, with +// and without ascale, and a K whose rotation block narrows. + +#include "foldquant.h" + +#include +#include +#include +#include +#include + +using namespace vla; + +namespace { + +// Natural-order Sylvester-Hadamard entry (i, j) of size n: (-1)^popcount(i & j). +float hadamard(int i, int j) { + return (__builtin_popcount((unsigned) (i & j)) & 1) ? -1.0f : 1.0f; +} + +int run(int64_t K, int64_t N, int wbits, bool fold_before, bool with_ascale, std::mt19937 & rng) { + const int rb = fq_rot_block_for(K, 64); + const int qmax = wbits == 4 ? 7 : 127; + std::uniform_int_distribution qd(-qmax, qmax); + std::uniform_real_distribution sd(0.001f, 0.02f), ad(0.5f, 2.0f); + + std::vector q((size_t) (N * K)); + for (auto & v : q) v = (int8_t) qd(rng); + std::vector ws((size_t) N), as((size_t) K); + for (auto & v : ws) v = sd(rng); + for (auto & v : as) v = ad(rng); + + // Codes as the GGUF stores them. + const int64_t kp = fq_w_kpack(K, wbits); + std::vector codes((size_t) (N * kp)); + for (int64_t n = 0; n < N; ++n) + for (int64_t k = 0; k < K; ++k) { + const int8_t v = q[(size_t) (n * K + k)]; + if (wbits == 8) { + codes[(size_t) (n * kp + k)] = v; + } else { + uint8_t & b = (uint8_t &) codes[(size_t) (n * kp + k / 2)]; + b = (k % 2 == 0) ? (uint8_t) ((b & 0xF0) | (v & 0xF)) : (uint8_t) ((b & 0x0F) | ((v & 0xF) << 4)); + } + } + + std::vector got((size_t) (N * K)); + fq_dequant_rows(codes.data(), ws.data(), with_ascale ? as.data() : nullptr, K, N, wbits, rb, fold_before, got.data()); + + // Dense reference in double. + const double norm = 1.0 / std::sqrt((double) rb); + double max_err = 0.0, max_ref = 0.0; + for (int64_t n = 0; n < N; ++n) { + std::vector w((size_t) K); + for (int64_t k = 0; k < K; ++k) w[(size_t) k] = (double) q[(size_t) (n * K + k)] * ws[(size_t) n]; + if (with_ascale && !fold_before) + for (int64_t k = 0; k < K; ++k) w[(size_t) k] /= as[(size_t) k]; + std::vector r((size_t) K, 0.0); + for (int64_t j = 0; j < K; ++j) { // (w . R)[j] = sum_i w[i] R[i][j] + const int64_t blk = j / rb * rb; + double acc = 0.0; + for (int64_t i = blk; i < blk + rb; ++i) + acc += w[(size_t) i] * hadamard((int) (i - blk), (int) (j - blk)) * norm; + r[(size_t) j] = rb > 1 ? acc : w[(size_t) j]; + } + if (with_ascale && fold_before) + for (int64_t k = 0; k < K; ++k) r[(size_t) k] /= as[(size_t) k]; + for (int64_t k = 0; k < K; ++k) { + max_err = std::fmax(max_err, std::fabs(r[(size_t) k] - got[(size_t) (n * K + k)])); + max_ref = std::fmax(max_ref, std::fabs(r[(size_t) k])); + } + } + const bool ok = max_err <= 1e-5 * std::fmax(1.0, max_ref); + std::printf("%s K=%lld rb=%d W%d fold_%s ascale=%d: max|err| %.3g (max|w| %.3g)\n", ok ? "ok " : "FAIL", + (long long) K, rb, wbits, fold_before ? "before" : "after", (int) with_ascale, max_err, max_ref); + return ok ? 0 : 1; +} + +} // namespace + +int main() { + std::mt19937 rng(1234); + int fails = 0; + for (int64_t K : {128, 96}) // 96 narrows the nominal 64-block to 32 + for (int wbits : {8, 4}) + for (bool before : {true, false}) + for (bool asc : {true, false}) + fails += run(K, 64, wbits, before, asc, rng); + std::printf("%s\n", fails ? "test_foldquant_dequant: FAILED" : "test_foldquant_dequant: all ok"); + return fails ? 1 : 0; +} diff --git a/tests/test_foldquant_ov_op.cpp b/tests/test_foldquant_ov_op.cpp new file mode 100644 index 0000000..8afd6cd --- /dev/null +++ b/tests/test_foldquant_ov_op.cpp @@ -0,0 +1,194 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant's two custom nodes on ggml's OpenVINO backend (translated by +// src/openvino/foldquant_ov.cpp) against the CPU reference (foldquant_ref.h): +// the cases of test_foldquant_cpu_op.cpp, with the site's tensors in a weights +// buffer as a real model has them, on the device GGML_OPENVINO_DEVICE names. +// +// Codes are computed in a different float order than the reference (OpenVINO's +// own reductions), so a value sitting on a rounding boundary can land one code +// over; the output must match the reference wherever every code agrees and may +// differ by one activation step where one does not. Built only with -DGGML_OPENVINO=ON. + +#include "foldquant.h" +#include "foldquant_ref.h" +#include "layers/fq_linear.h" + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-openvino.h" + +#include +#include +#include +#include +#include + +namespace { + +struct Lcg { + uint32_t s; + explicit Lcg(uint32_t seed) : s(seed) {} + float next() { + s = s * 1664525u + 1013904223u; + return ((float) ((s >> 8) & 0xFFFFu) / 65535.0f) * 2.0f - 1.0f; + } +}; + +struct Case { + const char * name; + int wbits, abits, rot; + bool gamma, ascale, fold_before, bias; +}; + +constexpr int64_t K = 128, N = 64, T = 5; + +int run_case(ggml_backend_t backend, const Case & c) { + Lcg rng(0x5eed1234u); + std::vector hx((size_t) K * T), hw((size_t) K * N), hws(N), hb(N), has(K), hga(K); + for (auto & v : hx) v = rng.next() * 4.0f; + for (auto & v : hw) v = rng.next(); + for (auto & v : hws) v = 0.01f + 0.02f * std::fabs(rng.next()); + for (auto & v : hb) v = rng.next() * 0.5f; + for (auto & v : has) v = 0.5f + std::fabs(rng.next()); + for (auto & v : hga) v = 0.75f + 0.5f * std::fabs(rng.next()); + + const float wq = c.wbits == 4 ? 7.0f : 127.0f; + std::vector wcodes((size_t) K * N); + for (int64_t n = 0; n < N; ++n) { + float amax = 0.f; + for (int64_t k = 0; k < K; ++k) amax = std::fmax(amax, std::fabs(hw[(size_t) n * K + k])); + for (int64_t k = 0; k < K; ++k) + wcodes[(size_t) n * K + k] = (int8_t) std::nearbyint(hw[(size_t) n * K + k] / (amax / wq)); + } + const int64_t kpw = vla::fq_w_kpack(K, c.wbits); + std::vector wpacked((size_t) kpw * N); + if (c.wbits == 4) { + for (int64_t n = 0; n < N; ++n) + for (int64_t k = 0; k < K; k += 2) + wpacked[(size_t) n * kpw + k / 2] = + (int8_t) ((wcodes[(size_t) n * K + k] & 0xF) | ((wcodes[(size_t) n * K + k + 1] & 0xF) << 4)); + } else { + wpacked = wcodes; + } + + // The site's tensors live in a weights buffer, as a model's do: that is what + // makes the backend turn them into constants. + ggml_init_params wp = { ggml_tensor_overhead() * 8, nullptr, true }; + ggml_context * W = ggml_init(wp); + ggml_tensor * w = ggml_new_tensor_2d(W, GGML_TYPE_I8, kpw, N); + ggml_tensor * ws = ggml_new_tensor_1d(W, GGML_TYPE_F32, N); + ggml_tensor * b = c.bias ? ggml_new_tensor_1d(W, GGML_TYPE_F32, N) : nullptr; + ggml_tensor * as = c.ascale ? ggml_new_tensor_1d(W, GGML_TYPE_F32, K) : nullptr; + ggml_tensor * ga = c.gamma ? ggml_new_tensor_1d(W, GGML_TYPE_F32, K) : nullptr; + ggml_set_name(w, "site.weight"); ggml_set_name(ws, "site.wscale"); + if (b) ggml_set_name(b, "site.bias"); + if (as) ggml_set_name(as, "site.ascale"); + if (ga) ggml_set_name(ga, "site.gamma"); + ggml_backend_buffer_t wbuf = ggml_backend_alloc_ctx_tensors(W, backend); + ggml_backend_buffer_set_usage(wbuf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_tensor_set(w, wpacked.data(), 0, ggml_nbytes(w)); + ggml_backend_tensor_set(ws, hws.data(), 0, ggml_nbytes(ws)); + if (b) ggml_backend_tensor_set(b, hb.data(), 0, ggml_nbytes(b)); + if (as) ggml_backend_tensor_set(as, has.data(), 0, ggml_nbytes(as)); + if (ga) ggml_backend_tensor_set(ga, hga.data(), 0, ggml_nbytes(ga)); + + vla::FqLinear s; + s.w = w; s.wscale = ws; s.bias = b; s.ascale = as; s.gamma = ga; + s.act.K = K; s.act.abits = c.abits; s.act.rot_block = c.rot; s.act.fold_before = c.fold_before; + s.act.has_gamma = c.gamma; s.act.has_ascale = c.ascale; s.act.clip = c.abits == 4 ? 0.9f : 1.0f; s.act.eps = 1e-6f; + s.gemm.K = K; s.gemm.N = N; s.gemm.wbits = c.wbits; + + ggml_init_params ip = { ggml_tensor_overhead() * 16 + ggml_graph_overhead(), nullptr, true }; + ggml_context * C = ggml_init(ip); + ggml_tensor * x = ggml_new_tensor_2d(C, GGML_TYPE_F32, K, T); + ggml_set_name(x, "x"); + ggml_set_input(x); + ggml_tensor * y = vla::fq_gemm(C, s, vla::fq_act(C, s, x)); + ggml_set_output(y); + ggml_cgraph * gf = ggml_new_graph(C); + ggml_build_forward_expand(gf, y); + ggml_gallocr_t ga_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(ga_alloc, gf)) { std::printf("FAIL: alloc\n"); return 1; } + ggml_backend_tensor_set(x, hx.data(), 0, ggml_nbytes(x)); + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + std::printf("FAIL: %s: graph compute\n", c.name); + return 1; + } + std::vector got((size_t) N * T); + ggml_backend_tensor_get(y, got.data(), 0, ggml_nbytes(y)); + + // CPU reference: the activation codes, then the integer GEMM. + int rows_exact = 0, rows_off = 0; + double max_rel = 0.0; + std::vector tmp(K); + float partial[vla::fqref::NT]; + const int64_t kp = vla::fq_act_kpack(K, c.abits); + std::vector blob((size_t) kp), codes((size_t) K); + for (int64_t t = 0; t < T; ++t) { + float xs; + vla::fqref::act_row(hx.data() + (size_t) t * K, as ? has.data() : nullptr, ga ? hga.data() : nullptr, + s.act, tmp.data(), partial, blob.data(), &xs); + if (c.abits == 4) vla::fqref::unpack_nibbles(blob.data(), K, codes.data()); + else codes = blob; + double row_err = 0.0, row_ref = 0.0; + for (int64_t n = 0; n < N; ++n) { + const float want = vla::fqref::gemm_dot(wcodes.data() + (size_t) n * K, codes.data(), K, xs, hws[n], + b ? hb[n] : 0.0f); + row_err = std::fmax(row_err, std::fabs((double) got[(size_t) t * N + n] - want)); + row_ref = std::fmax(row_ref, std::fabs((double) want)); + } + const double rel = row_err / std::fmax(row_ref, 1e-6); + max_rel = std::fmax(max_rel, rel); + // Every code equal: float accumulation differences only. One code over on a + // rounding boundary moves the row by at most |w| * xs * wscale per output. + if (rel <= 1e-4) ++rows_exact; else ++rows_off; + } + const bool ok = rows_off <= 1 && max_rel <= (c.abits == 4 ? 0.05 : 0.01); + std::printf("%s case %-20s W%dA%d rot%-2d gamma=%d ascale=%d before=%d bias=%d: %d/%lld rows exact, max rel %.2e\n", + ok ? "ok " : "FAIL", c.name, c.wbits, c.abits, c.rot, c.gamma, c.ascale, c.fold_before, c.bias, + rows_exact, (long long) T, max_rel); + ggml_gallocr_free(ga_alloc); + ggml_free(C); + ggml_backend_buffer_free(wbuf); + ggml_free(W); + return ok ? 0 : 1; +} + +} // namespace + +int main() { + ggml_backend_t backend = ggml_backend_openvino_init(0); + if (!backend) { + std::printf("SKIP: no OpenVINO backend\n"); + return 0; + } + const Case cases[] = { + { "w8a8_rot64", 8, 8, 64, false, false, false, true }, + { "w8a8_rot64_gamma", 8, 8, 64, true, false, false, false }, + { "w8a8_rot32_pre", 8, 8, 32, false, true, true, true }, + { "w8a8_rot64_post", 8, 8, 64, false, true, false, false }, + { "w8a8_norot", 8, 8, 1, false, false, false, false }, + { "w4a8_rot64_gamma", 4, 8, 64, true, false, false, true }, + { "w4a4_rot64_pre", 4, 4, 64, false, true, true, true }, + }; + int fails = 0; + for (const Case & c : cases) + fails += run_case(backend, c); + std::printf("%s\n", fails ? "test_foldquant_ov_op: FAILED" : "test_foldquant_ov_op: PASS"); + ggml_backend_free(backend); + return fails ? 1 : 0; +} diff --git a/tests/test_foldquant_sycl_op.cpp b/tests/test_foldquant_sycl_op.cpp new file mode 100644 index 0000000..6a51ec5 --- /dev/null +++ b/tests/test_foldquant_sycl_op.cpp @@ -0,0 +1,286 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FoldQuant's two custom nodes on ggml's SYCL backend (src/sycl/ +// vla_sycl_foldquant.cpp) against the CPU reference (foldquant_ref.h), on the +// cases of test_foldquant_cpu_op.cpp: the activation blob must match byte for +// byte and every output bit for bit, as the CUDA kernels do. Built only with +// -DGGML_SYCL=ON. +// +// With a GPU the nodes go through ggml's SYCL backend. ggml-sycl refuses to +// start without one, so otherwise they are handed straight to the extension hook +// on whatever SYCL device there is (the OpenCL CPU device on a CI runner, or +// ONEAPI_DEVICE_SELECTOR=opencl:cpu), with their tensors in USM shared memory: +// the same kernels, minus XMX and oneDNN where the device has neither. + +#include "foldquant.h" +#include "foldquant_ref.h" +#include "layers/fq_linear.h" +#include "sycl/vla_sycl_foldquant.h" + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-sycl.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" { +typedef bool (*ggml_sycl_ext_forward_t)(struct ggml_tensor * dst, void * queue); +extern ggml_sycl_ext_forward_t ggml_sycl_ext_forward; +} + +namespace { + +struct Lcg { + uint32_t s; + explicit Lcg(uint32_t seed) : s(seed) {} + float next() { + s = s * 1664525u + 1013904223u; + return ((float) ((s >> 8) & 0xFFFFu) / 65535.0f) * 2.0f - 1.0f; + } +}; + +struct Case { + const char * name; + int wbits, abits, rot; + bool gamma, ascale, fold_before, bias, residual; +}; + +// The shape under test; main() runs the cases over a list of them. +int64_t K = 256, N = 128, T = 37; + +// Exactly one of the two is set: ggml's SYCL backend, or a queue for the hook. +struct Exec { + ggml_backend_t backend = nullptr; + sycl::queue * q = nullptr; +}; + +int run_case(const Exec & ex, const Case & c) { + ggml_backend_t backend = ex.backend; + Lcg rng(0x5eed1234u); + std::vector hx((size_t) K * T), hw((size_t) K * N), hws(N), hb(N), has(K), hga(K), hr((size_t) N * T); + for (auto & v : hx) v = rng.next() * 4.0f; + for (auto & v : hw) v = rng.next(); + for (auto & v : hws) v = 0.01f + 0.02f * std::fabs(rng.next()); + for (auto & v : hb) v = rng.next() * 0.5f; + for (auto & v : has) v = 0.5f + std::fabs(rng.next()); + for (auto & v : hga) v = 0.75f + 0.5f * std::fabs(rng.next()); + for (auto & v : hr) v = rng.next(); + + const float wq = c.wbits == 4 ? 7.0f : 127.0f; + std::vector wcodes((size_t) K * N); + for (int64_t n = 0; n < N; ++n) { + float amax = 0.f; + for (int64_t k = 0; k < K; ++k) amax = std::fmax(amax, std::fabs(hw[(size_t) n * K + k])); + for (int64_t k = 0; k < K; ++k) + wcodes[(size_t) n * K + k] = (int8_t) std::nearbyint(hw[(size_t) n * K + k] / (amax / wq)); + } + const int64_t kpw = vla::fq_w_kpack(K, c.wbits); + std::vector wpacked((size_t) kpw * N); + if (c.wbits == 4) { + for (int64_t n = 0; n < N; ++n) + for (int64_t k = 0; k < K; k += 2) + wpacked[(size_t) n * kpw + k / 2] = + (int8_t) ((wcodes[(size_t) n * K + k] & 0xF) | ((wcodes[(size_t) n * K + k + 1] & 0xF) << 4)); + } else { + wpacked = wcodes; + } + + ggml_init_params wp = { ggml_tensor_overhead() * 8, nullptr, true }; + ggml_context * W = ggml_init(wp); + ggml_tensor * w = ggml_new_tensor_2d(W, GGML_TYPE_I8, kpw, N); + ggml_tensor * ws = ggml_new_tensor_1d(W, GGML_TYPE_F32, N); + ggml_tensor * b = c.bias ? ggml_new_tensor_1d(W, GGML_TYPE_F32, N) : nullptr; + ggml_tensor * as = c.ascale ? ggml_new_tensor_1d(W, GGML_TYPE_F32, K) : nullptr; + ggml_tensor * ga = c.gamma ? ggml_new_tensor_1d(W, GGML_TYPE_F32, K) : nullptr; + ggml_set_name(w, "site.weight"); + + vla::FqLinear s; + s.w = w; s.wscale = ws; s.bias = b; s.ascale = as; s.gamma = ga; + s.act.K = K; s.act.abits = c.abits; s.act.rot_block = c.rot; s.act.fold_before = c.fold_before; + s.act.has_gamma = c.gamma; s.act.has_ascale = c.ascale; s.act.clip = c.abits == 4 ? 0.9f : 1.0f; s.act.eps = 1e-6f; + s.gemm.K = K; s.gemm.N = N; s.gemm.wbits = c.wbits; + + ggml_init_params ip = { ggml_tensor_overhead() * 16 + ggml_graph_overhead(), nullptr, true }; + ggml_context * C = ggml_init(ip); + ggml_tensor * x = ggml_new_tensor_2d(C, GGML_TYPE_F32, K, T); + ggml_tensor * r = c.residual ? ggml_new_tensor_2d(C, GGML_TYPE_F32, N, T) : nullptr; + ggml_set_input(x); + if (r) ggml_set_input(r); + ggml_tensor * xq = vla::fq_act(C, s, x); + ggml_tensor * y = vla::fq_gemm(C, s, xq, r); + ggml_set_output(xq); + ggml_set_output(y); + ggml_cgraph * gf = ggml_new_graph(C); + ggml_build_forward_expand(gf, y); + if (ggml_graph_n_nodes(gf) != 2 || ggml_graph_node(gf, 0) != xq || ggml_graph_node(gf, 1) != y) { + std::printf("FAIL: %s: the graph is not fq_act -> fq_gemm\n", c.name); + return 1; + } + + const int64_t rb = vla::fq_act_row_bytes(K, c.abits), kp = vla::fq_act_kpack(K, c.abits); + std::vector got_blob((size_t) rb * T); + std::vector got_y((size_t) N * T); + const std::pair inputs[] = { + { w, wpacked.data() }, { ws, hws.data() }, { b, hb.data() }, { as, has.data() }, { ga, hga.data() }, + { x, hx.data() }, { r, hr.data() }, + }; + ggml_backend_buffer_t wbuf = nullptr; + ggml_gallocr_t alloc = nullptr; + std::vector usm; + if (backend) { + wbuf = ggml_backend_alloc_ctx_tensors(W, backend); + ggml_backend_buffer_set_usage(wbuf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(alloc, gf)) { std::printf("FAIL: alloc\n"); return 1; } + for (const auto & [t, h] : inputs) + if (t) ggml_backend_tensor_set(t, h, 0, ggml_nbytes(t)); + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + std::printf("FAIL: %s: graph compute\n", c.name); + return 1; + } + ggml_backend_tensor_get(xq, got_blob.data(), 0, got_blob.size()); + ggml_backend_tensor_get(y, got_y.data(), 0, got_y.size() * sizeof(float)); + } else { + auto place = [&](ggml_tensor * t, const void * h) { + void * p = sycl::malloc_shared(ggml_nbytes(t), *ex.q); + if (h) std::memcpy(p, h, ggml_nbytes(t)); + t->data = p; + usm.push_back(p); + }; + for (const auto & [t, h] : inputs) + if (t) place(t, h); + place(xq, nullptr); + place(y, nullptr); + // The two nodes in graph order, as ggml-sycl's compute loop would hand them over. + if (!ggml_sycl_ext_forward(xq, ex.q) || !ggml_sycl_ext_forward(y, ex.q)) { + std::printf("FAIL: %s: the hook declined a node\n", c.name); + return 1; + } + ex.q->wait(); + std::memcpy(got_blob.data(), xq->data, got_blob.size()); + std::memcpy(got_y.data(), y->data, got_y.size() * sizeof(float)); + } + + // Reference: blob rows, then the integer GEMM and its epilogue. + int bad_rows = 0, bad_y = 0; + std::vector tmp(K); + float partial[vla::fqref::NT]; + std::vector row((size_t) rb); + std::vector codes((size_t) K); + for (int64_t t = 0; t < T; ++t) { + std::memset(row.data(), 0, rb); + vla::fqref::act_row(hx.data() + (size_t) t * K, as ? has.data() : nullptr, ga ? hga.data() : nullptr, + s.act, tmp.data(), partial, (int8_t *) row.data(), (float *) (row.data() + kp)); + if (std::memcmp(row.data(), got_blob.data() + (size_t) t * rb, kp + 4) != 0) ++bad_rows; + if (c.abits == 4) vla::fqref::unpack_nibbles((const int8_t *) row.data(), K, codes.data()); + else std::memcpy(codes.data(), row.data(), K); + float xs; + std::memcpy(&xs, row.data() + kp, 4); + for (int64_t n = 0; n < N; ++n) { + float want = vla::fqref::gemm_dot(wcodes.data() + (size_t) n * K, codes.data(), K, xs, hws[n], + b ? hb[n] : 0.0f); + if (r) want = want + hr[(size_t) t * N + n]; + if (std::memcmp(&want, &got_y[(size_t) t * N + n], 4) != 0) { + if (bad_y < 3 && std::getenv("VLA_FQ_TEST_VERBOSE")) + std::printf(" y[%lld][%lld] got %.9g want %.9g (blob row %s)\n", (long long) t, (long long) n, + got_y[(size_t) t * N + n], want, + std::memcmp(row.data(), got_blob.data() + (size_t) t * rb, kp + 4) ? "differs" : "same"); + ++bad_y; + } + } + } + const bool ok = bad_rows == 0 && bad_y == 0; + std::printf("%s case %-22s W%dA%d rot%-2d gamma=%d ascale=%d before=%d bias=%d res=%d: " + "%d/%lld blob rows differ, %d/%lld outputs differ\n", + ok ? "ok " : "FAIL", c.name, c.wbits, c.abits, c.rot, c.gamma, c.ascale, c.fold_before, c.bias, + c.residual, bad_rows, (long long) T, bad_y, (long long) (N * T)); + for (void * p : usm) sycl::free(p, *ex.q); + if (alloc) ggml_gallocr_free(alloc); + ggml_free(C); + if (wbuf) ggml_backend_buffer_free(wbuf); + ggml_free(W); + return ok ? 0 : 1; +} + +} // namespace + +int main() { + if (sycl::device::get_devices().empty()) { + // CI sets VLA_FQ_TEST_REQUIRE_DEVICE so a runner that lost its OpenCL CPU + // device fails instead of passing on a skip. + const char * req = std::getenv("VLA_FQ_TEST_REQUIRE_DEVICE"); + std::printf("%s: no SYCL device\n", req && *req && *req != '0' ? "FAIL" : "SKIP"); + return req && *req && *req != '0' ? 1 : 0; + } + Exec ex; + std::unique_ptr q; + if (!sycl::device::get_devices(sycl::info::device_type::gpu).empty()) { + ex.backend = ggml_backend_sycl_init(0); + if (!ex.backend) { std::printf("FAIL: ggml_backend_sycl_init\n"); return 1; } + std::printf("through ggml's SYCL backend\n"); + } else { + q = std::make_unique(sycl::default_selector_v, sycl::property::queue::in_order()); + ex.q = q.get(); + std::printf("no GPU: through the extension hook on %s\n", + q->get_device().get_info().c_str()); + } + vla::sycl_register_foldquant_ops(); + const Case cases[] = { + { "w8a8_rot64", 8, 8, 64, false, false, false, true, false }, + { "w8a8_rot64_gamma", 8, 8, 64, true, false, false, false, true }, + { "w8a8_rot32_pre", 8, 8, 32, false, true, true, true, false }, + { "w8a8_rot64_post", 8, 8, 64, false, true, false, false, false }, + { "w8a8_norot", 8, 8, 1, false, false, false, false, false }, + { "w4a8_rot64_gamma", 4, 8, 64, true, false, false, true, true }, + { "w4a4_rot64_pre", 4, 4, 64, false, true, true, true, false }, + { "w8a4_rot64_gamma", 8, 4, 64, true, false, false, false, false }, + }; + // Small shapes over every path, then production ones: K = 16384 (a down + // projection: more chunks than the per-row act kernel holds, so the warp + // kernel), K = 6144 and 2048 rows through the per-row kernel, a prefill M + // through the XMX GEMM and a decode M through the GEMV, at full K so the int32 + // sums reach the magnitudes a real layer produces. + struct Shape { int64_t K, N, T; }; + const Shape shapes[] = { + { 256, 128, 37 }, { 256, 128, 10 }, { 256, 200, 37 }, + { 16384, 256, 300 }, { 16384, 256, 10 }, { 6144, 192, 41 }, { 2048, 2048, 64 }, + }; + // Every GEMM path: the default (oneDNN when built with it), this repo's + // GEMV + XMX kernels, and the plain tiled one. + int fails = 0; + for (const char * path : { "", "native", "simple" }) { + setenv("VLA_FQ_SYCL_GEMM", path, 1); + for (const Shape & sh : shapes) { + K = sh.K; N = sh.N; T = sh.T; + std::printf("GEMM path '%s', K = %lld, N = %lld, T = %lld\n", path, (long long) K, (long long) N, + (long long) T); + for (const Case & c : cases) + fails += run_case(ex, c); + } + } + std::printf("%s\n", fails ? "test_foldquant_sycl_op: FAILED" : "test_foldquant_sycl_op: PASS"); + if (ex.backend) ggml_backend_free(ex.backend); + return fails ? 1 : 0; +} diff --git a/tests/test_graph_names.cpp b/tests/test_graph_names.cpp index 0ee8d4f..06d1f25 100644 --- a/tests/test_graph_names.cpp +++ b/tests/test_graph_names.cpp @@ -20,7 +20,9 @@ // tensor's own name as the "%s" source, and glibc empties it instead of // appending, so every duplicate collapsed to "#". +#ifndef GGML_USE_OPENVINO // an OpenVINO build defines it already #define GGML_USE_OPENVINO +#endif #include "backend.h" #undef NDEBUG // keep assert() live even in Release builds