From 16c0ce1bb73a91c18fddd6459f8c17997c4fee71 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Wed, 30 Sep 2026 14:34:33 +0700 Subject: [PATCH 01/11] FoldQuant W8A8/W4A4 inference and FoldQuantVLA quantized-model converter Runtime - A FoldQuant GGUF carries the GR00T N1.5 / N1.6 / N1.7 and pi0.5 language backbone and 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. The format and arithmetic are specified in docs/QUANTIZATION.md. - Each site is two GGML_OP_CUSTOM nodes (src/layers/fq_linear.h). The CPU backend runs the reference in src/foldquant_ref.cpp; CUDA runs src/kernels/foldquant/: a fused RMSNorm + butterfly + quant prologue and an mma.sync INT8/INT4 GEMM with per-shape tiling, residual add and head layout in the epilogue. The CUDA extension ops go through one dispatcher (src/cuda/vla_cuda_ext.cu). Other backends refuse a FoldQuant file at load. - WeightLoader gains typed optional / fused declares; a float gemm() declare of an INT8 tensor fails with the FoldQuant site's name. - GR00T N1.7 precomputes the DiT adaLN conditions per denoising step at load. Converter - scripts/convert_quantized_model_to_gguf.py converts a FoldQuantVLA quantized model, either its quantized checkpoint (.qweight / .weight_scale in place of each quantized .weight) or its earlier fake-quant state, with no calibration and nothing re-rounded. A W4A4 DiT's INT4 adaLN is written dequantized; every quantized projection must become a GGUF site or that adaLN. --check-onnx byte-compares every site and folded norm gain against the TensorRT plugin graphs. - The family converters expose convert(ckpt, out, writer_factory=...); scripts/gguf_quant_writer.py is the quantizing writer. convert_pi05_to_gguf.py accepts an OpenPI checkpoint without config.json. - Tooling: foldquant_fake_export.py (uncalibrated file for bring-up), inspect_gguf_quant.py (contract check), foldquant_ref.py (numpy reference), foldquant_dequant.py (dequantizing exporter). Tests: CPU-op, CUDA-op and GEMM checks for FoldQuant; Python tests for the reference and the converter. --- CHANGELOG.md | 30 + CMakeLists.txt | 39 +- README.md | 22 + docs/ARCHITECTURE.md | 11 +- docs/QUANTIZATION.md | 234 ++++ scripts/convert_gr00t_n1_5_to_gguf.py | 59 +- scripts/convert_gr00t_n1_6_to_gguf.py | 59 +- scripts/convert_gr00t_n1_7_to_gguf.py | 26 +- scripts/convert_pi05_to_gguf.py | 83 +- scripts/convert_quantized_model_to_gguf.py | 1152 ++++++++++++++++++++ scripts/foldquant_dequant.py | 141 +++ scripts/foldquant_fake_export.py | 138 +++ scripts/foldquant_ref.py | 183 ++++ scripts/gguf_quant_writer.py | 139 +++ scripts/inspect_gguf_quant.py | 125 +++ scripts/quantize_gguf.py | 10 + src/backend.h | 6 +- src/cuda/vla_cuda_bf16.cu | 3 +- src/cuda/vla_cuda_ext.cu | 59 + src/cuda/vla_cuda_ext.h | 27 + src/cuda/vla_cuda_foldquant.cu | 312 ++++++ src/cuda/vla_cuda_ops.h | 11 +- src/foldquant.cpp | 324 ++++++ src/foldquant.h | 156 +++ src/foldquant_ref.cpp | 303 +++++ src/foldquant_ref.h | 87 ++ src/kernels/foldquant/fq_gemm_i8.cu | 392 +++++++ src/kernels/foldquant/fq_gemm_mma.cu | 324 ++++++ src/kernels/foldquant/fq_kernels.h | 97 ++ src/kernels/foldquant/fq_prologue.cu | 357 ++++++ src/layers/fq_linear.h | 95 ++ src/loader.cpp | 35 + src/loader.h | 9 + src/models/gr00tn1d5.cpp | 8 +- src/models/gr00tn1d6.cpp | 8 +- src/models/gr00tn1d7.cpp | 140 ++- src/models/pi05.cpp | 116 +- src/modules/dit_head.cpp | 175 ++- src/modules/dit_head.h | 29 +- src/modules/gemma_expert.h | 44 +- src/modules/qwen3_lm.cpp | 85 +- src/modules/qwen3_lm.h | 14 +- tests/CMakeLists.txt | 31 +- tests/foldquant_gemm_check.cu | 194 ++++ tests/py/test_foldquant_ref.py | 93 ++ tests/py/test_quantized_model_converter.py | 473 ++++++++ tests/test_foldquant_cpu_op.cpp | 294 +++++ tests/test_foldquant_cuda_op.cpp | 279 +++++ 48 files changed, 6820 insertions(+), 211 deletions(-) create mode 100644 docs/QUANTIZATION.md create mode 100644 scripts/convert_quantized_model_to_gguf.py create mode 100644 scripts/foldquant_dequant.py create mode 100644 scripts/foldquant_fake_export.py create mode 100644 scripts/foldquant_ref.py create mode 100644 scripts/gguf_quant_writer.py create mode 100644 scripts/inspect_gguf_quant.py create mode 100644 src/cuda/vla_cuda_ext.cu create mode 100644 src/cuda/vla_cuda_ext.h create mode 100644 src/cuda/vla_cuda_foldquant.cu create mode 100644 src/foldquant.cpp create mode 100644 src/foldquant.h create mode 100644 src/foldquant_ref.cpp create mode 100644 src/foldquant_ref.h create mode 100644 src/kernels/foldquant/fq_gemm_i8.cu create mode 100644 src/kernels/foldquant/fq_gemm_mma.cu create mode 100644 src/kernels/foldquant/fq_kernels.h create mode 100644 src/kernels/foldquant/fq_prologue.cu create mode 100644 src/layers/fq_linear.h create mode 100644 tests/foldquant_gemm_check.cu create mode 100644 tests/py/test_foldquant_ref.py create mode 100644 tests/py/test_quantized_model_converter.py create mode 100644 tests/test_foldquant_cpu_op.cpp create mode 100644 tests/test_foldquant_cuda_op.cpp diff --git a/CHANGELOG.md b/CHANGELOG.md index f98d4c6..0572933 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,36 @@ Notable changes to vla.cpp. Format loosely follows [Keep a Changelog](https://ke faster than BF16 on CPUs without BF16 matmul. - `VLA_BUILD_SERVER=OFF` builds `vla-cli` and `vla-bench` without protobuf or ZeroMQ. +- **FoldQuant INT8 / INT4 checkpoints.** A FoldQuant GGUF carries the GR00T + and π0.5 language backbone and action module as INT8 or INT4 codes in a block-Hadamard, SmoothQuant-folded frame with per-row scales; + vla.cpp quantizes activations per token and runs the projections on the + integer tensor cores. The file format and arithmetic are the contract in + `docs/QUANTIZATION.md`. Each site is two `GGML_OP_CUSTOM` nodes + (`src/layers/fq_linear.h`): the CPU backend runs the reference in + `src/foldquant_ref.cpp`, CUDA claims the same nodes through the ggml extension + hook with `src/kernels/foldquant/` (wmma INT8 GEMM, fused RMSNorm + butterfly + + quant prologue), bit-identical to the CPU path. `scripts/foldquant_fake_export.py` + produces an uncalibrated file for bring-up, `scripts/inspect_gguf_quant.py` + checks one against the contract, `scripts/foldquant_ref.py` is the numpy + reference. Other backends refuse a FoldQuant file at load. +- The ggml CUDA extension hook now goes through one dispatcher + (`src/cuda/vla_cuda_ext.cu`) so the BF16 activation ops and the FoldQuant + handler compose. +- `WeightLoader::opt_typed`, `fuse_typed`, `reader()` and `fail()`; `fuse()` + refuses sources of differing type or row shape, and a float `gemm()` declare + of an INT8 tensor fails with a message naming the FoldQuant site. +- The family converters expose `convert(ckpt, out, writer_factory=...)`, so a + quantizing writer reuses them instead of re-implementing the file. +- `scripts/convert_quantized_model_to_gguf.py` turns a quantized model into a + FoldQuant GGUF without calibration: FoldQuantVLA's quantized checkpoint + (`.qweight` / `.weight_scale` in place of each quantized `.weight`; the family + converters read it through a base view, and a W4A4 DiT's INT4 adaLN is + dequantized, since vla.cpp keeps adaLN in float), FoldQuantVLA's earlier + `foldquant.fakequant` state (sites by name, gains folded as the emitter + folds them). GR00T N1.5 / N1.6 / N1.7 and pi0.5; `--check-onnx` + byte-compares every site against the TensorRT plugin graphs. + `scripts/gguf_quant_writer.py` is the quantizing writer. + - **OpenVINO backend.** `-DGGML_OPENVINO=ON` runs the archs on Intel CPUs, iGPUs and NPUs through ggml's OpenVINO backend. SmolVLA, π0.5, Evo-1 and VLA-Adapter match an F32 CPU reference to 1e-3; on an Arc B390 iGPU that is 3.0x to 9.6x diff --git a/CMakeLists.txt b/CMakeLists.txt index 1066aea..8dabef3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -163,6 +163,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 @@ -189,6 +191,9 @@ 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 kernels bit for bit, so no +# FMA contraction on either side (aarch64 GCC contracts by default). +set_source_files_properties(src/foldquant.cpp src/foldquant_ref.cpp PROPERTIES COMPILE_OPTIONS "-ffp-contract=off") if(VLA_OCTO) target_sources(vla_core PRIVATE src/models/octo.cpp) target_include_directories(vla_core PRIVATE ${sentencepiece_SOURCE_DIR}/src) @@ -243,8 +248,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 @@ -258,7 +291,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() @@ -427,7 +460,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() foreach(tgt IN LISTS VLA_FIRST_PARTY_TARGETS) diff --git a/README.md b/README.md index 59a06af..2c4a6fa 100644 --- a/README.md +++ b/README.md @@ -255,6 +255,28 @@ python scripts/quantize_gguf.py --in model-bf16.gguf --out model-q8_0.gguf --typ 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 still runs float activations through `ggml_mul_mat`. 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, and vla.cpp quantizes +the activations per token and runs the GEMMs on the integer tensor cores. It loads like any other checkpoint on the CUDA and CPU backends (other backends +refuse it); the format and the arithmetic are in [docs/QUANTIZATION.md](docs/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 +``` + --- ## Benchmarks diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 0ed9960..3d65c42 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -20,6 +20,8 @@ source is the detail. - `src/serving/` - `vla-server` (ZeroMQ + protobuf, action prediction), `vlm-server` (chat), and `vla-cli` (one-shot inference). - `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 @@ -57,7 +59,14 @@ Two patterns, chosen per architecture: llama.cpp is fetched and pinned by CMake `FetchContent`; a bump is a one-line `GIT_TAG` change. Weights are bf16 by default and can be repacked to Q8_0/Q4_0 with `scripts/quantize_gguf.py`; the loader runs quantized GGUFs directly and lets -`ggml_mul_mat` dequantize at compute. CPU thread count scales to the machine core +`ggml_mul_mat` dequantize at compute. + +A FoldQuant GGUF (see [QUANTIZATION.md](QUANTIZATION.md)) +carries INT8 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/`). CPU thread count scales to the machine core count; CUDA and Metal run the towers and the transformer on the GPU. ## Adding an architecture diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md new file mode 100644 index 0000000..3898360 --- /dev/null +++ b/docs/QUANTIZATION.md @@ -0,0 +1,234 @@ +# 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 | CUDA, CPU (others refuse the file at load) | + +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. Other backends +refuse the file at load: there is no per-op fallback in vla.cpp. + +Environment switches: `VLA_FQ_CHECK=1` 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/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 e8e9280..fc4c514 100644 --- a/scripts/convert_pi05_to_gguf.py +++ b/scripts/convert_pi05_to_gguf.py @@ -121,6 +121,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 { @@ -149,30 +152,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 MEAN_STD 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} cfg = dict(GEMMA_2B, **GEMMA_300M) cfg["paligemma_variant"] = str(cfg_json.get("paligemma_variant", "gemma_2b")) @@ -239,15 +244,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(PFX_VLM_HEAD)) @@ -266,10 +271,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 MEAN_STD 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/quantize_gguf.py b/scripts/quantize_gguf.py index 6e37647..3dc6a59 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 codes + sidecar scales executed by +the in-tree integer kernels, CUDA/CPU only) 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 @@ -87,6 +94,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) meta = {"GGUF.version", "GGUF.tensor_count", "GGUF.kv_count", "general.architecture"} diff --git a/src/backend.h b/src/backend.h index 087638e..19625dc 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 70a3d63..dc943fc 100644 --- a/src/cuda/vla_cuda_bf16.cu +++ b/src/cuda/vla_cuda_bf16.cu @@ -37,6 +37,7 @@ #include "ggml.h" #include "env_flag.h" +#include "cuda/vla_cuda_ext.h" #include #include @@ -723,7 +724,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..a7e34b8 --- /dev/null +++ b/src/foldquant.cpp @@ -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. + +#include "foldquant.h" + +#include "backend.h" +#include "env_flag.h" +#include "gguf_reader.h" +#include "loader.h" +#include "cuda/vla_cuda_ops.h" + +#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, const 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; + if (!b.is_cuda && !is_cpu) { + std::fprintf(stderr, + "%s: a FoldQuant GGUF runs on the CUDA or CPU backend only (this build drives '%s'). " + "Rebuild with -DGGML_CUDA=ON or without an accelerator, or use the bf16 GGUF.\n", + tag, name); + return false; + } + 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; +} + +namespace { + +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; + 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 + + 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) 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"); + } + + 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..2aa713e --- /dev/null +++ b/src/foldquant.h @@ -0,0 +1,156 @@ +// 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. +// +// Both backends run the same two GGML_OP_CUSTOM nodes per site (fq_act, +// fq_gemm; see layers/fq_linear.h): the CPU backend executes the reference in +// foldquant_ref.h, the CUDA backend claims them through the ggml extension hook. + +#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; + + 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; + + 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 (registers the kernels) or CPU. Everything else is +// refused: GGML_OP_CUSTOM has no implementation there and the core drives one +// backend with no per-op fallback. +bool foldquant_check_backend(const char * tag, const Backend & b, const FoldQuantSpec & fq, bool weight_dtype_set); + +// 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. `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 0137731..aa51689 100644 --- a/src/loader.cpp +++ b/src/loader.cpp @@ -47,6 +47,16 @@ 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) { + 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 +95,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); } @@ -115,6 +142,14 @@ ggml_tensor * WeightLoader::fuse(ggml_type want, const char * out_name, const st ok_ = false; return nullptr; } + // Concatenation along the last axis needs the same row shape and, for a + // packed type copied raw, the same source type. + if (gs->type != first->type || (!is1d && gs->ne[0] != first->ne[0]) || ggml_n_dims(gs) != ggml_n_dims(first)) { + std::fprintf(stderr, "vla(%s): cannot fuse %s with %s (type/shape differ)\n", + arch_, s.c_str(), srcs[0].c_str()); + ok_ = false; + return nullptr; + } rows += is1d ? gs->ne[0] : gs->ne[1]; } diff --git a/src/loader.h b/src/loader.h index 649c9e9..de04665 100644 --- a/src/loader.h +++ b/src/loader.h @@ -43,6 +43,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 +56,14 @@ 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); + + // 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_; diff --git a/src/models/gr00tn1d5.cpp b/src/models/gr00tn1d5.cpp index 03a70ea..0a1424e 100644 --- a/src/models/gr00tn1d5.cpp +++ b/src/models/gr00tn1d5.cpp @@ -16,6 +16,7 @@ #include "options.h" #include "backend.h" #include "env_flag.h" +#include "foldquant.h" #include "gguf_reader.h" #include "layers/embed.h" #include "layers/linear.h" @@ -232,10 +233,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"); + const 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); @@ -250,7 +254,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"); @@ -258,7 +262,7 @@ std::unique_ptr gr00t_n1_5_create(const std::string& mmproj_path, m->aex.declare(L, "aex"); 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 c017f36..8e8748e 100644 --- a/src/models/gr00tn1d6.cpp +++ b/src/models/gr00tn1d6.cpp @@ -16,6 +16,7 @@ #include "options.h" #include "backend.h" #include "env_flag.h" +#include "foldquant.h" #include "gguf_reader.h" #include "layers/embed.h" #include "layers/ffn.h" @@ -266,10 +267,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"); + const 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); @@ -289,13 +293,13 @@ 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"); m->aex.declare(L, "aex"); - 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 0e959b3..4636de5 100644 --- a/src/models/gr00tn1d7.cpp +++ b/src/models/gr00tn1d7.cpp @@ -26,6 +26,7 @@ #include "ggml.h" #include "ggml-cpu.h" #include "ggml-backend.h" +#include "ggml-alloc.h" #include "backend.h" #include "gguf.h" #include "gguf_reader.h" @@ -66,6 +67,15 @@ struct Gr00tN1d7ModelArch : public ModelArchBase { scratch_ctx vision_scratch; ggml_backend_buffer_t weight_buf = nullptr; ggml_type matmul_type = GGML_TYPE_F32; + // adaLN conditions per (denoising step, DiT layer) and the proj_out condition + // per step, computed once at load: the timestep embedding is fixed per step, + // so these GEMVs (2*hidden x hidden, bf16, 132 per request) never change. + // VLA_GR00T_ADALN_CACHE=0 rebuilds them in the graph instead (bit-identical). + ggml_context * ctx_cache = nullptr; + ggml_backend_buffer_t cache_buf = nullptr; + std::vector t_cond; // [s*dit_layers + i], F32[2*dit_hidden] + std::vector t_po; // [s], F32[2*dit_hidden] + bool precompute_dit_cond(); int64_t vit_hidden=1024, vit_layers=24, vit_heads=16, vit_inter=4096; int64_t patch_size=16, temporal_patch=2, spatial_merge=2, vit_num_pos=2304, vit_patch_flat=1536, vit_merged_dim=4096; @@ -228,6 +238,10 @@ bool load_config(const gguf_reader & g, Gr00tN1d7ModelArch & m, Config & cfg) { Gr00tN1d7ModelArch::~Gr00tN1d7ModelArch() { mg.release(); + if (cache_buf) + ggml_backend_buffer_free(cache_buf); + if (ctx_cache) + ggml_free(ctx_cache); if (weight_buf) ggml_backend_buffer_free(weight_buf); if (ctx_weights) @@ -256,6 +270,7 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, } if (!load_config(g, *m, m->cfg)) return nullptr; + const 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", @@ -273,6 +288,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 }; @@ -285,14 +302,25 @@ 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->vit_layers); - 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"); m->vlsa.declare(L, "aex.vlsa", m->vlsa_layers, EncNames{"norm1", "norm3", "ff0", "ff2"}); m->aex.declare(L, "aex"); - 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; @@ -305,9 +333,69 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, std::fprintf(stderr, "vla(gr00tn1d7): build_caches failed\n"); return nullptr; } + if (env_flag("VLA_GR00T_ADALN_CACHE", true) && !m->precompute_dit_cond()) { + std::fprintf(stderr, "vla(gr00tn1d7): adaLN condition precompute failed\n"); + return nullptr; + } return m; } +// Runs the timestep MLP and every layer's adaLN GEMV once per denoising step +// on the backend (the same ops the graph would run) and keeps the results +// resident, so predict() reads them as leaves. +bool Gr00tN1d7ModelArch::precompute_dit_cond() { + const int64_t L = dit_layers, S = num_steps, D2 = 2*dit_hidden; + ggml_init_params p = { (size_t) 16*1024*1024, nullptr, true }; + ggml_context * C = ggml_init(p); + if (!C) return false; + ggml_cgraph * gf = ggml_new_graph_custom(C, 4096, false); + std::vector tproj((size_t) S), conds((size_t) (S*L)), pos((size_t) S); + for (int64_t s=0; s> h_cond((size_t) (S*L), std::vector((size_t) D2)), h_po((size_t) S, std::vector((size_t) D2)); + for (int64_t k=0; k Gr00tN1d7ModelArch::predict(const Inputs& in) { ? ggml_new_tensor_1d(C, GGML_TYPE_I32, SEQ_TXT) : nullptr; if (t_txt_idx) ggml_set_input(t_txt_idx); - std::vector t_tau(num_steps), t_tproj(num_steps); + const bool cond_cached = !t_cond.empty(); + std::vector t_tau(num_steps), t_tproj(num_steps, nullptr); for (int64_t s=0; s Gr00tN1d7ModelArch::predict(const Inputs& in) { ggml_tensor * actions = t_x0; for (int64_t s=0; snb[1], 0)); @@ -588,9 +678,11 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { else if (i%every2 == 0) enc = vl_txt; else enc = vl_img; - hh = dit.block(C, dit.blk[i], hh, temb, enc, Kc[i], Vc[i]); + hh = dit.block(C, dit.blk[i], hh, temb, enc, Kc[i], Vc[i], + cond_cached ? t_cond[(size_t) (s*dit_layers + i)] : nullptr); } - ggml_tensor * po = ggml_add(C, ggml_mul_mat(C, dit.po1W, ggml_silu(C, temb)), dit.po1b); + ggml_tensor * po = cond_cached ? t_po[(size_t) s] + : ggml_add(C, ggml_mul_mat(C, dit.po1W, ggml_silu(C, temb)), dit.po1b); ggml_tensor * sh = ggml_view_1d(C, po, dit_hidden, 0), * sc = ggml_view_1d(C, po, dit_hidden, (size_t) dit_hidden * sizeof(float)); ggml_tensor * hn = ggml_norm(C, hh, norm_out_eps); ggml_tensor * h_mod = ggml_add(C, ggml_add(C, hn, ggml_mul(C, hn, sc)), sh); @@ -697,10 +789,42 @@ std::vector Gr00tN1d7ModelArch::predict(const Inputs& in) { ggml_backend_tensor_set(t_txt_idx, text_pos_idx.data(), 0, ggml_nbytes(t_txt_idx)); for (int64_t s=0; s 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 02a5cce..5490e4f 100644 --- a/src/models/pi05.cpp +++ b/src/models/pi05.cpp @@ -29,6 +29,8 @@ #include "layers/embed.h" #include "modules/preprocess.h" #include "env_flag.h" +#include "foldquant.h" +#include "layers/fq_linear.h" #include #include @@ -60,6 +62,9 @@ struct ExpertLayerW { ggml_tensor * Wgate = nullptr; ggml_tensor * Wup = nullptr; ggml_tensor * Wdown = nullptr; + // FoldQuant sites (action recipe): the adaRMS output is quantized per site + // group with the shipped SmoothQuant vector (ascale), no gamma. + FqLinear fq_q, fq_k, fq_v, fq_o, fq_gate, fq_up, fq_down; }; bool ends_with(const std::string & s, const char * sfx) { @@ -159,11 +164,20 @@ ggml_tensor * build_vlm_layer( const int64_t nkv = cfg.n_kv_heads; const int64_t qf = nq * hd; - ggml_tensor * x_norm = ggml_mul(ctx, ggml_rms_norm(ctx, x_in, cfg.rms_eps), w.ln_in); - - ggml_tensor * q = ggml_mul_mat(ctx, w.Wq, x_norm); - ggml_tensor * k = ggml_mul_mat(ctx, w.Wk, x_norm); - ggml_tensor * v = ggml_mul_mat(ctx, w.Wv, x_norm); + ggml_tensor *q, *k, *v; + if (w.fq_q) { + // FoldQuant: RMSNorm with the folded gamma is fused into the act node, + // which feeds all three projections. + ggml_tensor * xq = fq_act(ctx, w.fq_q, x_in); + q = fq_gemm(ctx, w.fq_q, xq); + k = fq_gemm(ctx, w.fq_k, xq); + v = fq_gemm(ctx, w.fq_v, xq); + } else { + ggml_tensor * x_norm = ggml_mul(ctx, ggml_rms_norm(ctx, x_in, cfg.rms_eps), w.ln_in); + q = ggml_mul_mat(ctx, w.Wq, x_norm); + k = ggml_mul_mat(ctx, w.Wk, x_norm); + v = ggml_mul_mat(ctx, w.Wv, x_norm); + } ggml_tensor * q_h = ggml_reshape_3d(ctx, q, hd, nq, seq); ggml_tensor * k_h = ggml_reshape_3d(ctx, k, hd, nkv, seq); @@ -194,15 +208,23 @@ ggml_tensor * build_vlm_layer( ggml_tensor * att_pre = ggml_reshape_2d(ctx, ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)), qf, seq); - ggml_tensor * o_out = ggml_mul_mat(ctx, w.Wo, att_pre); - ggml_tensor * h1 = ggml_add(ctx, x_in, o_out); - - ggml_tensor * x_norm_mlp = ggml_mul(ctx, ggml_rms_norm(ctx, h1, cfg.rms_eps), w.ln_post); - ggml_tensor * gate = ggml_mul_mat(ctx, w.Wgate, x_norm_mlp); - ggml_tensor * up = ggml_mul_mat(ctx, w.Wup, x_norm_mlp); + // FoldQuant: the residual adds ride in the o / down GEMM epilogues. + ggml_tensor * h1 = w.fq_o ? fq_linear(ctx, w.fq_o, att_pre, x_in) + : ggml_add(ctx, x_in, ggml_mul_mat(ctx, w.Wo, att_pre)); + + ggml_tensor *gate, *up; + if (w.fq_gate) { + ggml_tensor * xq2 = fq_act(ctx, w.fq_gate, h1); + gate = fq_gemm(ctx, w.fq_gate, xq2); + up = fq_gemm(ctx, w.fq_up, xq2); + } else { + ggml_tensor * x_norm_mlp = ggml_mul(ctx, ggml_rms_norm(ctx, h1, cfg.rms_eps), w.ln_post); + gate = ggml_mul_mat(ctx, w.Wgate, x_norm_mlp); + up = ggml_mul_mat(ctx, w.Wup, x_norm_mlp); + } ggml_tensor * inter_t = ggml_mul(ctx, vla::gelu(ctx, gate), up); - ggml_tensor * mlp_out = ggml_mul_mat(ctx, w.Wdown, inter_t); - return ggml_add(ctx, h1, mlp_out); + return w.fq_down ? fq_linear(ctx, w.fq_down, inter_t, h1) + : ggml_add(ctx, h1, ggml_mul_mat(ctx, w.Wdown, inter_t)); } ggml_tensor * build_adarms( @@ -236,9 +258,17 @@ ggml_tensor * build_expert_layer( ggml_tensor * gate_attn = nullptr; ggml_tensor * x_norm = build_adarms(ctx, x_in, w.ada_in_w, w.ada_in_b, cond, h, cfg.rms_eps, &gate_attn); - ggml_tensor * q = ggml_mul_mat(ctx, w.Wq, x_norm); - ggml_tensor * k = ggml_mul_mat(ctx, w.Wk, x_norm); - ggml_tensor * v = ggml_mul_mat(ctx, w.Wv, x_norm); + ggml_tensor *q, *k, *v; + if (w.fq_q) { + ggml_tensor * xq = fq_act(ctx, w.fq_q, x_norm); // q/k/v share the act (same ascale) + q = fq_gemm(ctx, w.fq_q, xq); + k = fq_gemm(ctx, w.fq_k, xq); + v = fq_gemm(ctx, w.fq_v, xq); + } else { + q = ggml_mul_mat(ctx, w.Wq, x_norm); + k = ggml_mul_mat(ctx, w.Wk, x_norm); + v = ggml_mul_mat(ctx, w.Wv, x_norm); + } ggml_tensor * q_h = ggml_reshape_3d(ctx, q, hd, nq, seq); ggml_tensor * k_h = ggml_reshape_3d(ctx, k, hd, nkv, seq); @@ -267,16 +297,24 @@ ggml_tensor * build_expert_layer( ggml_tensor * att_pre = ggml_reshape_2d(ctx, ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)), qf, seq); - ggml_tensor * o_out = ggml_mul_mat(ctx, w.Wo, att_pre); + // The expert's residuals are gated (out * gate + x), so no epilogue fusion here. + ggml_tensor * o_out = w.fq_o ? fq_linear(ctx, w.fq_o, att_pre) : ggml_mul_mat(ctx, w.Wo, att_pre); 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, w.ada_post_w, w.ada_post_b, cond, h, cfg.rms_eps, &gate_ffn); - ggml_tensor * gate = ggml_mul_mat(ctx, w.Wgate, x_norm_mlp); - ggml_tensor * up = ggml_mul_mat(ctx, w.Wup, x_norm_mlp); + ggml_tensor *gate, *up; + if (w.fq_gate) { + ggml_tensor * xq2 = fq_act(ctx, w.fq_gate, x_norm_mlp); + gate = fq_gemm(ctx, w.fq_gate, xq2); + up = fq_gemm(ctx, w.fq_up, xq2); + } else { + gate = ggml_mul_mat(ctx, w.Wgate, x_norm_mlp); + up = ggml_mul_mat(ctx, w.Wup, x_norm_mlp); + } ggml_tensor * inter_t = ggml_mul(ctx, vla::gelu(ctx, gate), up); - ggml_tensor * mlp_out = ggml_mul_mat(ctx, w.Wdown, inter_t); + ggml_tensor * mlp_out = w.fq_down ? fq_linear(ctx, w.fq_down, inter_t) : ggml_mul_mat(ctx, w.Wdown, inter_t); return ggml_add(ctx, h1, ggml_mul(ctx, mlp_out, gate_ffn)); } @@ -404,6 +442,7 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, if (!m->io.open(ckpt_path)) return nullptr; gguf_reader & g = m->io; + const FoldQuantSpec fq = foldquant_parse(g, "pi05"); if (!g.has("pi05.architecture") || g.str("pi05.architecture") != "pi05") { std::fprintf(stderr, "vla(pi05): '%s' is not a π0.5 GGUF (pi05.architecture missing/wrong)\n", ckpt_path.c_str()); @@ -431,6 +470,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. @@ -464,22 +505,41 @@ std::unique_ptr pi05_create(const std::string& mmproj_path, m->mm_proj_w = L.gemm ("mm.proj.weight"); m->mm_proj_b = L.opt_f32("mm.proj.bias"); - 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.ada_in_w = L.f32 ("aex.blk.%lld.attn_norm.weight", (long long)i); w.ada_in_b = L.f32 ("aex.blk.%lld.attn_norm.bias", (long long)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.ada_post_w = L.f32 ("aex.blk.%lld.ffn_norm.weight", (long long)i); w.ada_post_b = L.f32 ("aex.blk.%lld.ffn_norm.bias", (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); + if (fq.present) { + const FqModuleSpec & a = fq.action; + const long long ii = (long long) i; + 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", (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); + } + if (!w.fq_o) + w.Wo = L.gemm("aex.blk.%lld.attn_o.weight", (long long)i); + if (!w.fq_gate) { + 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); + } + if (!w.fq_down) + w.Wdown = L.gemm("aex.blk.%lld.ffn_down.weight", (long long)i); } m->ex_final_w = L.f32("aex.output_norm.weight"); diff --git a/src/modules/dit_head.cpp b/src/modules/dit_head.cpp index f081288..7c9aaff 100644 --- a/src/modules/dit_head.cpp +++ b/src/modules/dit_head.cpp @@ -16,6 +16,7 @@ #include "layers/attn.h" #include "layers/ffn.h" +#include "layers/fq_linear.h" #include "layers/linear.h" #include "layers/norm.h" @@ -24,7 +25,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; @@ -36,22 +38,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); } +// adaLN from a precomputed condition: the tail of layers/norm.h adaln, same ops. +static ggml_tensor * adaln_from_cond(ggml_context * C, ggml_tensor * x, ggml_tensor * cond, int64_t dim, float eps) { + ggml_tensor * sc = ggml_view_1d(C, cond, dim, 0); + ggml_tensor * sh = ggml_view_1d(C, cond, dim, (size_t) dim * sizeof(float)); + ggml_tensor * xn = ggml_norm(C, x, eps); + return ggml_add(C, ggml_add(C, xn, ggml_mul(C, xn, sc)), sh); +} + +ggml_tensor * DitHead::adaln_cond(ggml_context * C, const DitLayerW & w, ggml_tensor * temb) const { + return linear(C, w.adaln_w, w.adaln_b, ggml_silu(C, temb)); +} + +ggml_tensor * DitHead::proj_out_cond(ggml_context * C, ggml_tensor * temb) const { + return linear(C, po1W, po1b, ggml_silu(C, temb)); +} + ggml_tensor * DitHead::block(ggml_context * C, const DitLayerW & w, ggml_tensor * h, ggml_tensor * temb, - ggml_tensor * enc, ggml_tensor * K_pre, ggml_tensor * V_pre) const { + ggml_tensor * enc, ggml_tensor * K_pre, ggml_tensor * V_pre, ggml_tensor * cond) const { const int64_t hd = cfg.head_dim; const int64_t heads = cfg.heads; const int64_t dim = cfg.hidden; const int64_t Tk = h->ne[1]; const float scale = 1.0f/std::sqrt((float)hd); - ggml_tensor * n = adaln(C, h, temb, w.adaln_w, w.adaln_b, dim, cfg.ln_eps); + ggml_tensor * n = cond ? adaln_from_cond(C, h, cond, dim, cfg.ln_eps) + : adaln(C, h, temb, w.adaln_w, w.adaln_b, dim, cfg.ln_eps); 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)); } @@ -136,8 +237,8 @@ ggml_tensor * DitHead::time_emb(ggml_context * C, ggml_tensor * tproj) const { return linear(C, te_l2W, te_l2b, ggml_silu(C, linear(C, te_l1W, te_l1b, tproj))); } -ggml_tensor * DitHead::proj_out(ggml_context * C, ggml_tensor * h, ggml_tensor * temb) const { - ggml_tensor * po = linear(C, po1W, po1b, ggml_silu(C, temb)); +ggml_tensor * DitHead::proj_out(ggml_context * C, ggml_tensor * h, ggml_tensor * temb, ggml_tensor * po) const { + if (!po) po = proj_out_cond(C, temb); ggml_tensor * sh = ggml_view_1d(C, po, cfg.hidden, 0); ggml_tensor * sc = ggml_view_1d(C, po, cfg.hidden, (size_t)cfg.hidden*sizeof(float)); diff --git a/src/modules/dit_head.h b/src/modules/dit_head.h index 6c00b2c..2b2a249 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" @@ -33,6 +34,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 { @@ -52,19 +57,31 @@ 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; + // cond: this layer's adaLN condition (scale, shift) [2*hidden] already + // computed from temb (see adaln_cond); temb is then unused. ggml_tensor * block(ggml_context * C, const DitLayerW & w, ggml_tensor * h, ggml_tensor * temb, - ggml_tensor * enc, ggml_tensor * K_pre = nullptr, ggml_tensor * V_pre = nullptr) const; + ggml_tensor * enc, ggml_tensor * K_pre = nullptr, ggml_tensor * V_pre = nullptr, + ggml_tensor * cond = nullptr) const; ggml_tensor * time_emb(ggml_context * C, ggml_tensor * tproj) const; - - // (shift, scale) adaLN, opposite to layers/norm.h adaln. - ggml_tensor * proj_out(ggml_context * C, ggml_tensor * h, ggml_tensor * temb) const; + // The adaLN condition of one layer, adaln_w . silu(temb) + adaln_b: it depends + // on the timestep embedding only, so a model whose temb is fixed per + // denoising step computes it once at load and passes it to block() as cond, + // saving one 2*hidden x hidden GEMV per layer per step per request. + ggml_tensor * adaln_cond(ggml_context * C, const DitLayerW & w, ggml_tensor * temb) const; + // po1 . silu(temb) + po1_b, the (shift, scale) of proj_out; same reuse as adaln_cond. + ggml_tensor * proj_out_cond(ggml_context * C, ggml_tensor * temb) const; + + // (shift, scale) adaLN, opposite to layers/norm.h adaln. po: a precomputed proj_out_cond. + ggml_tensor * proj_out(ggml_context * C, ggml_tensor * h, ggml_tensor * temb, ggml_tensor * po = nullptr) const; }; } diff --git a/src/modules/gemma_expert.h b/src/modules/gemma_expert.h index 33e1001..2af12f7 100644 --- a/src/modules/gemma_expert.h +++ b/src/modules/gemma_expert.h @@ -18,6 +18,7 @@ #pragma once +#include "foldquant.h" #include "loader.h" #include "ggml.h" @@ -37,25 +38,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/tests/CMakeLists.txt b/tests/CMakeLists.txt index ddffcfd..9c9fce4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -52,6 +52,14 @@ target_link_libraries(test_config_guard PRIVATE vla_core) target_compile_options(test_config_guard PRIVATE -Wall -Wextra) add_test(NAME config_guard COMMAND test_config_guard) +# 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) + # 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. @@ -61,12 +69,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) endif() 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; +} From 430e9f1bdbcc19119d77baa6dfe4af07829e8387 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 2 Oct 2026 12:50:04 +0700 Subject: [PATCH 02/11] FoldQuant: read sites back as float weights on backends without the kernels Metal, SYCL, OpenVINO, Hexagon and OpenCL have no implementation of the two FoldQuant custom nodes and vla.cpp drives one backend with no per-op fallback, so a FoldQuant GGUF was refused there. Now foldquant_check_backend switches such a backend (or CUDA/CPU with VLA_FQ_DEQUANT=1) to dequant mode: every site is registered with the loader as a float GEMM weight, rebuilt at upload in the resident type, and the arch takes its stock float path. The activation path is x' = R(x / a) (fold before) or R(x) / a (after), R the block-normalised Sylvester-Hadamard butterfly (symmetric, orthonormal), so 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. - WeightLoader::as_float registers a tensor of another file type as a float [K, N] weight; gemm/opt_gemm, fuse_gemm (DiT fused q/k/v, k/v) and upload honour it, converting to the resident type. - The four arch callers hold the spec mutable so the check can set the mode. - test_foldquant_dequant checks the rebuilt weights against the dense product (W8/W4, both fold orders, with and without ascale, a narrowed block). Weight-only quantization on those backends: weights keep FoldQuant's rounding, activations stay float. pi0.5 LIBERO W4A4 (vrfai/pi05-libero-w4a4): action cosine vs bf16 0.99986 dequant, 0.99703 integer path; dequant CUDA vs CPU 1.00000. GR00T N1.7 LIBERO W4A4: integer vs dequant 0.99930, dequant CUDA vs CPU 1.00000. --- CHANGELOG.md | 11 ++-- docs/ARCHITECTURE.md | 3 +- docs/MODELS.md | 6 +- docs/QUANTIZATION.md | 28 +++++++- scripts/quantize_gguf.py | 4 +- src/foldquant.cpp | 98 +++++++++++++++++++++++++-- src/foldquant.h | 30 +++++++-- src/loader.cpp | 69 ++++++++++++++++--- src/loader.h | 18 +++++ src/models/gr00tn1d5.cpp | 2 +- src/models/gr00tn1d6.cpp | 2 +- src/models/gr00tn1d7.cpp | 2 +- src/models/pi05.cpp | 2 +- tests/CMakeLists.txt | 8 +++ tests/test_foldquant_dequant.cpp | 110 +++++++++++++++++++++++++++++++ 15 files changed, 354 insertions(+), 39 deletions(-) create mode 100644 tests/test_foldquant_dequant.cpp diff --git a/CHANGELOG.md b/CHANGELOG.md index 918379f..2de81ef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,13 +6,16 @@ Notable changes to vla.cpp. Format loosely follows [Keep a Changelog](https://ke ### Added -- **FoldQuant W8A8 / W4A4 inference** for GR00T N1.5 / N1.6 / N1.7 and π0.5 on - CUDA and CPU. A FoldQuant GGUF carries the language backbone and the action +- **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, the - reference in `src/foldquant_ref.cpp` on CPU). Other backends refuse the file - at load. The format and arithmetic are in `docs/QUANTIZATION.md`. + reference in `src/foldquant_ref.cpp` on CPU). On every other backend (Metal, + Vulkan, SYCL, OpenVINO, Hexagon, OpenCL) 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 diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index f99f988..67a22da 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -79,7 +79,8 @@ carries INT8 or INT4 codes plus sidecar scales instead. `src/foldquant.h` declar 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/`). CPU thread count scales to the machine core count; +extension hook (`src/cuda/`). 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 2dd434f..eb7b427 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -50,8 +50,10 @@ 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 the CUDA and CPU backends (other backends -refuse it); the format and the arithmetic are in [QUANTIZATION.md](QUANTIZATION.md). +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 runs the integer kernels +and CPU the exact reference, and 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 diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md index 3898360..7a4cd15 100644 --- a/docs/QUANTIZATION.md +++ b/docs/QUANTIZATION.md @@ -8,7 +8,7 @@ vla.cpp runs two kinds of quantized checkpoint: | 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 | CUDA, CPU (others refuse the file at load) | +| Backends | all | integer kernels on CUDA, exact reference on CPU; 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 @@ -155,8 +155,30 @@ 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. Other backends -refuse the file at load: there is no per-op fallback in vla.cpp. +on the unsupported op rather than computing something else. + +### Other backends + +Metal, Vulkan, SYCL, OpenVINO, Hexagon and OpenCL 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` 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 diff --git a/scripts/quantize_gguf.py b/scripts/quantize_gguf.py index b36ce06..f364d92 100644 --- a/scripts/quantize_gguf.py +++ b/scripts/quantize_gguf.py @@ -24,8 +24,8 @@ 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 codes + sidecar scales executed by -the in-tree integer kernels, CUDA/CPU only) is produced from a FoldQuantVLA +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. """ diff --git a/src/foldquant.cpp b/src/foldquant.cpp index a7e34b8..03ee311 100644 --- a/src/foldquant.cpp +++ b/src/foldquant.cpp @@ -19,7 +19,9 @@ #include "gguf_reader.h" #include "loader.h" #include "cuda/vla_cuda_ops.h" +#include "foldquant_ref.h" +#include #include #include #include @@ -120,17 +122,20 @@ FoldQuantSpec foldquant_parse(const gguf_reader & g, const char * prefix) { return fq; } -bool foldquant_check_backend(const char * tag, const Backend & b, const FoldQuantSpec & fq, bool weight_dtype_set) { +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; - if (!b.is_cuda && !is_cpu) { - std::fprintf(stderr, - "%s: a FoldQuant GGUF runs on the CUDA or CPU backend only (this build drives '%s'). " - "Rebuild with -DGGML_CUDA=ON or without an accelerator, or use the bf16 GGUF.\n", - tag, name); - return false; + 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", @@ -151,8 +156,76 @@ bool foldquant_check_backend(const char * tag, const Backend & b, const FoldQuan 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); @@ -226,6 +299,10 @@ FqLinear fq_declare_linear(WeightLoader & L, const FqModuleSpec & mod, const cha 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); @@ -281,6 +358,13 @@ FqLinear fq_declare_fused(WeightLoader & L, const FqModuleSpec & mod, const char 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); diff --git a/src/foldquant.h b/src/foldquant.h index 2aa713e..8411be3 100644 --- a/src/foldquant.h +++ b/src/foldquant.h @@ -20,9 +20,17 @@ // the activation is [RMS-normed,] [divided,] rotated by the same butterfly, // quantized per token to INT8, multiplied on integer units and dequantized. // -// Both backends run the same two GGML_OP_CUSTOM nodes per site (fq_act, +// CUDA and CPU run the same two GGML_OP_CUSTOM nodes per site (fq_act, // fq_gemm; see layers/fq_linear.h): the CPU backend executes the reference in // foldquant_ref.h, the CUDA backend claims them through the ggml extension hook. +// +// Every other backend (Metal, Vulkan, SYCL, OpenVINO, Hexagon, OpenCL) has 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 @@ -56,6 +64,7 @@ struct FqModuleSpec { 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) int wbits_for(const char * site_key) const; int abits_for(const char * site_key) const; @@ -127,14 +136,23 @@ int fq_rot_block_for(int64_t K, int nominal); bool foldquant_present(const gguf_reader & g, const char * prefix); FoldQuantSpec foldquant_parse (const gguf_reader & g, const char * prefix); -// Load-time policy: CUDA (registers the kernels) or CPU. Everything else is -// refused: GGML_OP_CUSTOM has no implementation there and the core drives one -// backend with no per-op fallback. -bool foldquant_check_backend(const char * tag, const Backend & b, const FoldQuantSpec & fq, bool weight_dtype_set); +// Load-time policy: CUDA registers the integer kernels and CPU runs the exact +// reference; any other backend switches fq to dequant mode, as does +// VLA_FQ_DEQUANT=1 on CUDA or CPU. 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. `gamma` is the already +// 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, diff --git a/src/loader.cpp b/src/loader.cpp index 40d4df8..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, @@ -51,6 +78,18 @@ ggml_tensor * WeightLoader::declare(ggml_type want, bool required, bool gemma_no // 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; @@ -133,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()); @@ -145,17 +193,18 @@ ggml_tensor * WeightLoader::fuse(ggml_type want, const char * out_name, const st } // Same resident type and row shape; a tensor copied raw (a packed type, // FoldQuant INT8 codes) also needs every source in that same type. - if (g_.resident_type(gs, want) != rt || (rt == first->type && gs->type != first->type) || - (!is1d && gs->ne[0] != first->ne[0]) || ggml_n_dims(gs) != ggml_n_dims(first)) { + 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; @@ -188,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); @@ -201,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 de04665..87d09d0 100644 --- a/src/loader.h +++ b/src/loader.h @@ -24,6 +24,9 @@ #include "ggml-backend.h" #include +#include +#include +#include #include #include @@ -58,6 +61,12 @@ class WeightLoader { 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_; @@ -77,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 2eaa186..719f261 100644 --- a/src/models/gr00tn1d5.cpp +++ b/src/models/gr00tn1d5.cpp @@ -227,7 +227,7 @@ 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"); - const FoldQuantSpec fq = foldquant_parse(g, "gr00t_n1_5"); + FoldQuantSpec fq = foldquant_parse(g, "gr00t_n1_5"); const Backend b = backend_init("vla(gr00tn1d5)", m->n_threads); if (!b.handle) return nullptr; diff --git a/src/models/gr00tn1d6.cpp b/src/models/gr00tn1d6.cpp index 0ebd4bb..5e9b131 100644 --- a/src/models/gr00tn1d6.cpp +++ b/src/models/gr00tn1d6.cpp @@ -244,7 +244,7 @@ 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"); - const FoldQuantSpec fq = foldquant_parse(g, "gr00t_n1_6"); + FoldQuantSpec fq = foldquant_parse(g, "gr00t_n1_6"); const Backend b = backend_init("vla(gr00tn1d6)", m->n_threads); if (!b.handle) return nullptr; diff --git a/src/models/gr00tn1d7.cpp b/src/models/gr00tn1d7.cpp index 71e2e11..d0ce433 100644 --- a/src/models/gr00tn1d7.cpp +++ b/src/models/gr00tn1d7.cpp @@ -197,7 +197,7 @@ std::unique_ptr gr00t_n1_7_create(const std::string& mmproj_path, } if (!load_config(g, opts, *m, m->cfg)) return nullptr; - const FoldQuantSpec fq = foldquant_parse(g, "gr00t_n1_7"); + 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", diff --git a/src/models/pi05.cpp b/src/models/pi05.cpp index f0792ca..5992b0c 100644 --- a/src/models/pi05.cpp +++ b/src/models/pi05.cpp @@ -251,7 +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; - const FoldQuantSpec fq = foldquant_parse(g, "pi05"); + 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 " diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 9f13dcc..c72db68 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -69,6 +69,14 @@ 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) + # 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. 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; +} From d3f97144b346396e5a02925b1bfe769511f6b9ae Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 2 Oct 2026 14:42:32 +0700 Subject: [PATCH 03/11] FoldQuant on OpenVINO: translate the custom nodes into OpenVINO ops ggml's OpenVINO backend now runs a FoldQuant GGUF natively on the CPU and GPU plugins instead of reading the sites back as float weights. - src/openvino/foldquant_ov.cpp, compiled into ggml-openvino: fq_act becomes the CPU 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, per-token amax scale, round-half-even, clamp) and passes codes and scale as floats to fq_gemm, which multiplies them with the weight kept as an i8 / i4 constant (W4 bytes reinterpreted in place), dequantized by wscale in the decompression pattern the plugins keep compressed. - patch_ggml_openvino.py hunk 14: register the translator for GGML_OP_CUSTOM, let supports_op accept FoldQuant's nodes (only those) ahead of the type checks, keep INT8 weights as i8 constants. - foldquant_check_backend: native on OpenVINO CPU/GPU, dequant on the NPU or with VLA_FQ_DEQUANT=1; the head-laid-out epilogue is off there (FqModuleSpec::no_heads), since a translated graph has no raw layout. - test_foldquant_ov_op: one site through the backend against the CPU reference (W8A8, W4A8, W4A4; gamma, ascale before/after, bias, narrowed and no rotation): every code agrees, outputs within 1e-6 relative, on an Intel Core Ultra X7 358H CPU and its Arc B390 iGPU. Whole model, against the CUDA integer path: W8A8 pi0.5 1.00000 (CPU), 0.99999 (GPU). W4A4 pi0.5 / GR00T N1.7 0.9993 / 0.9987 (CPU): with 15 activation levels a code on a rounding tie flips by one when the reductions run in a different float order, and the flips compound. pi0.5 W4A4 on the CPU plugin: 3926 ms vs 6595 ms bf16, 6.7 GB vs 16.5 GB peak RSS; on the iGPU 651 ms vs 560 ms, 3.8 GB vs 11.6 GB. --- CHANGELOG.md | 7 +- CMakeLists.txt | 13 ++ docs/ARCHITECTURE.md | 5 +- docs/MODELS.md | 7 +- docs/QUANTIZATION.md | 33 ++++- docs/backend/ov.md | 25 ++++ scripts/patch_ggml_openvino.py | 65 ++++++++- src/foldquant.cpp | 21 ++- src/foldquant.h | 9 +- src/openvino/foldquant_ov.cpp | 235 +++++++++++++++++++++++++++++++++ tests/CMakeLists.txt | 9 ++ tests/test_foldquant_ov_op.cpp | 194 +++++++++++++++++++++++++++ 12 files changed, 607 insertions(+), 16 deletions(-) create mode 100644 src/openvino/foldquant_ov.cpp create mode 100644 tests/test_foldquant_ov_op.cpp diff --git a/CHANGELOG.md b/CHANGELOG.md index 2de81ef..97c8ffb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,9 +11,10 @@ Notable changes to vla.cpp. Format loosely follows [Keep a Changelog](https://ke 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, the - reference in `src/foldquant_ref.cpp` on CPU). On every other backend (Metal, - Vulkan, SYCL, OpenVINO, Hexagon, OpenCL) the sites are read back as float - weights with FoldQuant's rounding and run the stock float path; + reference in `src/foldquant_ref.cpp` on CPU, OpenVINO ops on the OpenVINO CPU + and GPU plugins with INT8/INT4 weight constants). On every other backend + (Metal, SYCL, 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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 1198fbd..3a8af90 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -117,6 +117,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 @@ -340,6 +351,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 diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 67a22da..7f7939b 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -79,8 +79,9 @@ carries INT8 or INT4 codes plus sidecar scales instead. `src/foldquant.h` declar 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/`). 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; +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 eb7b427..8776ab0 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -50,9 +50,10 @@ 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 runs the integer kernels -and CPU the exact reference, and the others read the sites back as float weights -with FoldQuant's rounding (weight-only quantization, bf16-GGUF speed); the format +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 runs the 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 diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md index 7a4cd15..6b11d21 100644 --- a/docs/QUANTIZATION.md +++ b/docs/QUANTIZATION.md @@ -8,7 +8,7 @@ vla.cpp runs two kinds of quantized checkpoint: | 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; every other backend reads the sites back as float weights (see [Other backends](#other-backends)) | +| 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 @@ -157,10 +157,37 @@ claimed by the extension hook (`src/cuda/vla_cuda_foldquant.cu`, registered by 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. + ### Other backends -Metal, Vulkan, SYCL, OpenVINO, Hexagon and OpenCL have no implementation of the -two custom nodes, and vla.cpp drives a single backend with no per-op fallback. +Metal, Vulkan, SYCL, 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 diff --git a/docs/backend/ov.md b/docs/backend/ov.md index 2393349..36b4b94 100644 --- a/docs/backend/ov.md +++ b/docs/backend/ov.md @@ -287,6 +287,31 @@ 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 | 6595 | 16.5 GB | +| CPU | FoldQuant W4A4 | 3926 | 6.7 GB | +| CPU | FoldQuant W8A8 (uncalibrated) | 3804 | 9.4 GB | +| GPU | bf16 | 560 | 11.6 GB | +| GPU | FoldQuant W4A4 | 651 | 3.8 GB | +| GPU | FoldQuant W8A8 (uncalibrated) | 607 | 9.8 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. + ## Known issues **π0 needs F32 on the GPU, and gets it by default.** The GPU plugin computes in 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/src/foldquant.cpp b/src/foldquant.cpp index 03ee311..abb58ed 100644 --- a/src/foldquant.cpp +++ b/src/foldquant.cpp @@ -127,6 +127,24 @@ bool foldquant_check_backend(const char * tag, const Backend & b, FoldQuantSpec return true; const char * name = b.handle ? ggml_backend_name(b.handle) : ""; const bool is_cpu = std::strcmp(name, "CPU") == 0; + 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. @@ -277,6 +295,7 @@ bool fill_specs(WeightLoader & L, const FqModuleSpec & mod, const char * site_ke r.gemm.K = K; r.gemm.N = N; r.gemm.wbits = wbits; + r.heads_off = mod.no_heads; return true; } @@ -319,7 +338,7 @@ FqLinear fq_declare_linear(WeightLoader & L, const FqModuleSpec & mod, const cha 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) return; + 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; diff --git a/src/foldquant.h b/src/foldquant.h index 8411be3..53ef5e9 100644 --- a/src/foldquant.h +++ b/src/foldquant.h @@ -65,6 +65,7 @@ struct FqModuleSpec { 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; @@ -118,6 +119,7 @@ struct FqLinear { 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; } }; @@ -136,9 +138,10 @@ int fq_rot_block_for(int64_t K, int nominal); bool foldquant_present(const gguf_reader & g, const char * prefix); FoldQuantSpec foldquant_parse (const gguf_reader & g, const char * prefix); -// Load-time policy: CUDA registers the integer kernels and CPU runs the exact -// reference; any other backend switches fq to dequant mode, as does -// VLA_FQ_DEQUANT=1 on CUDA or CPU. Returns false only on an unusable spec. +// Load-time policy: CUDA registers the 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 diff --git a/src/openvino/foldquant_ov.cpp b/src/openvino/foldquant_ov.cpp new file mode 100644 index 0000000..af53ec4 --- /dev/null +++ b/src/openvino/foldquant_ov.cpp @@ -0,0 +1,235 @@ +// 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+1] F32: the K integer-valued codes, +// then 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 +#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: FoldQuant's nodes, and no other GGML_OP_CUSTOM. +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); + auto out = std::make_shared(OutputVector{q, scale}, -1); + 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); + + // Activation codes and per-token scale from fq_act. + const Output act = context.get_input(1); + auto codes = std::make_shared(act, i64_vec({0}), i64_vec({K}), i64_vec({1}), i64_vec({-1})); + auto tok_s = std::make_shared(act, i64_vec({K}), i64_vec({K + 1}), i64_vec({1}), i64_vec({-1})); + + Output y = std::make_shared(codes, wf, false, true); + y = std::make_shared(y, tok_s); + + // 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/tests/CMakeLists.txt b/tests/CMakeLists.txt index c72db68..0af195f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -77,6 +77,15 @@ 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 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. 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; +} From b37a4fc82252b1cfa371c87160a982b1c608b68f Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 2 Oct 2026 16:22:50 +0700 Subject: [PATCH 04/11] FoldQuant on OpenVINO: carry the activation as codes times scale fq_act now emits the integer-valued codes already multiplied by their per-token scale, and fq_gemm feeds that straight into the MatMul with the decompressed weight. The scale factors out of the GEMM, so this drops the Concat, the two Slices and the scale Multiply at every site with the same arithmetic. Measured on an Intel Core Ultra X7 358H, pi0.5 LIBERO W4A4, against the codes+scale form: iGPU 666 -> 645 ms, CPU 4.69 -> 4.15 s (CPU timings on that machine drift by ~20% between runs). Accuracy against the CUDA integer path is unchanged: W8A8 1.00000 (CPU) / 0.99999 (GPU), W4A4 + o/down INT8 0.99962 / 0.99971, W4A4 0.99936 / 0.99928. test_foldquant_ov_op passes on the Intel CPU, the Arc B390 iGPU and OpenVINO's ARM CPU plugin. Tried and dropped: the plugins' dynamic activation quantization (on by default on the CPU, group 32, and the fastest setting there; no effect on the iGPU), and a group-quantized [N, K/g, g] weight layout (no faster on the iGPU). The iGPU stays behind bf16 because its compressed-weight matmul decompresses to F16 rather than running int8. --- src/openvino/foldquant_ov.cpp | 25 +++++++++++-------------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/src/openvino/foldquant_ov.cpp b/src/openvino/foldquant_ov.cpp index af53ec4..14a4c4d 100644 --- a/src/openvino/foldquant_ov.cpp +++ b/src/openvino/foldquant_ov.cpp @@ -18,9 +18,9 @@ // points (the op table entry and supports_op) are added by // scripts/patch_ggml_openvino.py. // -// fq_act (x[, gamma][, ascale]) -> [.., T, K+1] F32: the K integer-valued codes, -// then 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_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 @@ -44,7 +44,6 @@ #include #include #include -#include #include #include #include @@ -56,7 +55,6 @@ #include #include #include -#include #include #include @@ -78,7 +76,9 @@ uint32_t fq_magic(const void * userdata) { } // namespace -// supports_op hook: FoldQuant's nodes, and no other GGML_OP_CUSTOM. +// 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) { @@ -170,7 +170,9 @@ OutputVector translate_act(const NodeContext & context, const vla::FqActSpec & s std::make_shared(std::make_shared(y, inv), ov::op::v5::Round::RoundMode::HALF_TO_EVEN), -qmax, qmax); - auto out = std::make_shared(OutputVector{q, scale}, -1); + // 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()); } @@ -192,13 +194,8 @@ OutputVector translate_gemm(const NodeContext & context, const vla::FqGemmSpec & 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); - // Activation codes and per-token scale from fq_act. - const Output act = context.get_input(1); - auto codes = std::make_shared(act, i64_vec({0}), i64_vec({K}), i64_vec({1}), i64_vec({-1})); - auto tok_s = std::make_shared(act, i64_vec({K}), i64_vec({K + 1}), i64_vec({1}), i64_vec({-1})); - - Output y = std::make_shared(codes, wf, false, true); - y = std::make_shared(y, tok_s); + // 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. From fd2b8f52d4475ce8436e7575cbde14b75960f501 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 2 Oct 2026 18:09:06 +0700 Subject: [PATCH 05/11] FoldQuant on SYCL: bit-exact integer kernels behind a ggml-sycl extension hook pi0.5 W4A4 on an Arc B390 iGPU runs in 291 ms (bf16: 326 ms), W4A4 with INT8 o/down in 290 ms, W8A8 in 330 ms; every node bit-identical to the CPU reference. - scripts/patch_ggml_sycl_ext_hook.py: the SYCL counterpart of the CUDA hook, two exported pointers (compute, supports_op for GGML_OP_CUSTOM) consulted first by ggml-sycl; null by default, so a hooked ggml behaves as stock. - src/sycl/vla_sycl_foldquant.cpp: - fq_act follows the reference's reduction tree (lane l of a 32-wide sub-group owns chunks l, l+32, ...). Rows up to K = 8192 spread over a work-group (the CUDA act_row_kernel, ported); longer ones run one sub-group per row. - fq_gemm runs oneDNN's int8 matmul (INT4 weights as s4, INT4 activations unpacked to s8) into an int32 buffer, then the reference epilogue. Integer sums are exact in any order. Without oneDNN, or with VLA_FQ_SYCL_GEMM=native, a GEMV (M <= 32) and an int8 joint_matrix (XMX) kernel take over; =simple keeps a plain tiled one. - Exactness: -ffp-contract=off for the source, and the device JIT option -cl-fp32-correctly-rounded-divide-sqrt (a link option); without it the driver's division and sqrt move codes across rounding ties. The CPU reference gets -fp-model=precise under icpx, whose host default (-fp-model=fast) reassociates its sums. - foldquant_check_backend registers the kernels on a SYCL backend (plain [N][T] outputs, so no head layout); VLA_FQ_DEQUANT=1 still selects dequant. - test_foldquant_sycl_op: W8A8/W4A8/W4A4/W8A4 with gamma, ascale before/after, bias, residual, narrowed and no rotation, over every GEMM path and production shapes (K = 16384 / 6144 / 2048, M up to 300): 168 cases, activation bytes and outputs bit-identical. - Docs: QUANTIZATION.md SYCL section, backend/sycl.md numbers and the Level Zero loader workaround, MODELS.md, CHANGELOG, the foldquant.h overview. Whole model vs the CUDA integer path: W8A8 1.00000 action cosine, W4A4 + o/down INT8 0.99966, W4A4 0.99912 (pi0.5) / 0.9993 (GR00T N1.7); the gap is the float layers between the sites landing codes on other rounding ties. --- CHANGELOG.md | 10 +- CMakeLists.txt | 36 +- docs/MODELS.md | 6 +- docs/QUANTIZATION.md | 29 +- docs/backend/sycl.md | 35 ++ scripts/patch_ggml_sycl_ext_hook.py | 104 ++++ src/foldquant.cpp | 17 + src/foldquant.h | 18 +- src/sycl/vla_sycl_foldquant.cpp | 779 ++++++++++++++++++++++++++++ src/sycl/vla_sycl_foldquant.h | 25 + tests/CMakeLists.txt | 12 + tests/test_foldquant_sycl_op.cpp | 221 ++++++++ 12 files changed, 1273 insertions(+), 19 deletions(-) create mode 100644 scripts/patch_ggml_sycl_ext_hook.py create mode 100644 src/sycl/vla_sycl_foldquant.cpp create mode 100644 src/sycl/vla_sycl_foldquant.h create mode 100644 tests/test_foldquant_sycl_op.cpp diff --git a/CHANGELOG.md b/CHANGELOG.md index 97c8ffb..07b6029 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,10 +10,12 @@ Notable changes to vla.cpp. Format loosely follows [Keep a Changelog](https://ke 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, the - reference in `src/foldquant_ref.cpp` on CPU, OpenVINO ops on the OpenVINO CPU - and GPU plugins with INT8/INT4 weight constants). On every other backend - (Metal, SYCL, Hexagon, OpenCL, the OpenVINO NPU) the sites are read back as + 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`. diff --git a/CMakeLists.txt b/CMakeLists.txt index 3a8af90..f87a439 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} @@ -244,9 +250,12 @@ 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 kernels bit for bit, so no -# FMA contraction on either side (aarch64 GCC contracts by default). -set_source_files_properties(src/foldquant.cpp src/foldquant_ref.cpp PROPERTIES COMPILE_OPTIONS "-ffp-contract=off") +# 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), and no +# reassociation either: oneAPI's icpx compiles the host with -fp-model=fast by +# default, which reorders the reference's sums and products. +set_source_files_properties(src/foldquant.cpp src/foldquant_ref.cpp PROPERTIES COMPILE_OPTIONS + "$<$:-fp-model=precise>;$<$:-Wno-overriding-option>;-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}) @@ -340,6 +349,27 @@ 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. + 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;-fp-model=precise;-Wno-overriding-option;-ffp-contract=off") + target_link_options(vla_core PUBLIC -fsycl + "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. diff --git a/docs/MODELS.md b/docs/MODELS.md index 8776ab0..75b8b3c 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -50,9 +50,9 @@ 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 runs the 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 +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). diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md index 6b11d21..161f3d0 100644 --- a/docs/QUANTIZATION.md +++ b/docs/QUANTIZATION.md @@ -183,9 +183,36 @@ 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. A whole π0.5 model follows the CUDA integer +path to 1.00000 action cosine at W8A8, 0.9997 at W4A4 with INT8 o/down and 0.9991 +at W4A4. 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, SYCL, Hexagon and OpenCL (and the OpenVINO NPU) have no +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 diff --git a/docs/backend/sycl.md b/docs/backend/sycl.md index 01adcd5..d97b7d6 100644 --- a/docs/backend/sycl.md +++ b/docs/backend/sycl.md @@ -126,6 +126,41 @@ 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 (uncalibrated) | 4.73 GB | 330 | + +Vision (79 ms) stays bf16 in every row. W4A4 halves the weights and is 11% faster +than bf16; W8A8 runs at bf16 speed. + +## 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/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/src/foldquant.cpp b/src/foldquant.cpp index abb58ed..4b7c9df 100644 --- a/src/foldquant.cpp +++ b/src/foldquant.cpp @@ -20,6 +20,9 @@ #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 @@ -127,6 +130,20 @@ bool foldquant_check_backend(const char * tag, const Backend & b, FoldQuantSpec 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 diff --git a/src/foldquant.h b/src/foldquant.h index 53ef5e9..28a87ba 100644 --- a/src/foldquant.h +++ b/src/foldquant.h @@ -20,11 +20,13 @@ // the activation is [RMS-normed,] [divided,] rotated by the same butterfly, // quantized per token to INT8, multiplied on integer units and dequantized. // -// CUDA and CPU run the same two GGML_OP_CUSTOM nodes per site (fq_act, -// fq_gemm; see layers/fq_linear.h): the CPU backend executes the reference in -// foldquant_ref.h, the CUDA backend claims them through the ggml extension hook. +// 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). // -// Every other backend (Metal, Vulkan, SYCL, OpenVINO, Hexagon, OpenCL) has no +// 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 @@ -138,10 +140,10 @@ int fq_rot_block_for(int64_t K, int nominal); bool foldquant_present(const gguf_reader & g, const char * prefix); FoldQuantSpec foldquant_parse (const gguf_reader & g, const char * prefix); -// Load-time policy: CUDA registers the 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. +// 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 diff --git a/src/sycl/vla_sycl_foldquant.cpp b/src/sycl/vla_sycl_foldquant.cpp new file mode 100644 index 0000000..a4b7926 --- /dev/null +++ b/src/sycl/vla_sycl_foldquant.cpp @@ -0,0 +1,779 @@ +// 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 "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(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; +} + +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); + if (m == FQ_ACT_MAGIC) + return run_act(q, dst, *static_cast(ud)); + const FqGemmSpec & s = *static_cast(ud); + return run_gemm(q, dst, s, blob_abits(dst->src[1], s.K)); +} + +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 0af195f..3869632 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -77,6 +77,18 @@ 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: precise floating point, no contraction. + target_compile_options(test_foldquant_sycl_op PRIVATE -Wall -Wextra -fp-model=precise -Wno-overriding-option + -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) diff --git a/tests/test_foldquant_sycl_op.cpp b/tests/test_foldquant_sycl_op.cpp new file mode 100644 index 0000000..51ef686 --- /dev/null +++ b/tests/test_foldquant_sycl_op.cpp @@ -0,0 +1,221 @@ +// 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. + +#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 + +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; + +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), 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"); + 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_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); + ggml_gallocr_t 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; } + ggml_backend_tensor_set(x, hx.data(), 0, ggml_nbytes(x)); + if (r) ggml_backend_tensor_set(r, hr.data(), 0, ggml_nbytes(r)); + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + std::printf("FAIL: %s: graph compute\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); + 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)); + + // 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)); + ggml_gallocr_free(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_sycl_init(0); + if (!backend) { + std::printf("SKIP: no SYCL device\n"); + return 0; + } + 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(backend, c); + } + } + std::printf("%s\n", fails ? "test_foldquant_sycl_op: FAILED" : "test_foldquant_sycl_op: PASS"); + ggml_backend_free(backend); + return fails ? 1 : 0; +} From c04fce285591427e23561fc95cdd0727951206fd Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 2 Oct 2026 19:05:00 +0700 Subject: [PATCH 06/11] docs: OpenVINO FoldQuant numbers re-measured on a quiet machine, pointer to SYCL --- docs/backend/ov.md | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/docs/backend/ov.md b/docs/backend/ov.md index 36b4b94..bb38f97 100644 --- a/docs/backend/ov.md +++ b/docs/backend/ov.md @@ -301,16 +301,20 @@ B390 iGPU, OpenVINO 2026.4, `vla-bench` p50 over 10 calls after 3 warmups: | Device | Checkpoint | p50 ms | Peak RSS | |---|---|--:|--:| -| CPU | bf16 | 6595 | 16.5 GB | -| CPU | FoldQuant W4A4 | 3926 | 6.7 GB | -| CPU | FoldQuant W8A8 (uncalibrated) | 3804 | 9.4 GB | -| GPU | bf16 | 560 | 11.6 GB | -| GPU | FoldQuant W4A4 | 651 | 3.8 GB | -| GPU | FoldQuant W8A8 (uncalibrated) | 607 | 9.8 GB | +| 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 (uncalibrated) | 3763 | 9.4 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 (uncalibrated) | 595 | 9.6 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. +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)). ## Known issues From 48fca62786799b5e93544a832e1bf36ae5d67665 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Sat, 3 Oct 2026 11:40:14 +0700 Subject: [PATCH 07/11] CI: build and test the SYCL and OpenVINO backends - build-sycl: oneAPI 2026.1 (DPC++, oneMKL, oneDNN) on ubuntu-24.04, full build with -Werror, then ctest on the OpenCL CPU device the compiler package ships. ggml-sycl will not start without a GPU, so test_foldquant_sycl_op now hands its two nodes straight to the extension hook on any SYCL device, tensors in USM shared memory, and checks them against the CPU reference as before (GPU: still through ggml-sycl). VLA_FQ_TEST_REQUIRE_DEVICE turns a missing device into a failure. - build-openvino: OpenVINO 2026.4 archive (cached), full build with -Werror, ctest; the CPU plugin runs test_foldquant_ov_op on the runner. - build-gate checks the SYCL hook patch still applies, and is idempotent. - icpx compiles first-party code with -fp-model=precise (its default is fast, which broke test_vision_common's exact float compare); ggml keeps its own flags. The intended -ffp-contract=off override is no longer warned about. - The oneDNN GEMM path falls back to the native kernels, once per queue, on a device oneDNN cannot drive. --- .github/workflows/build.yml | 85 +++++++++++++++++++++++ CMakeLists.txt | 19 ++++-- docs/QUANTIZATION.md | 5 +- src/sycl/vla_sycl_foldquant.cpp | 16 +++++ tests/CMakeLists.txt | 6 +- tests/test_foldquant_sycl_op.cpp | 113 ++++++++++++++++++++++++------- 6 files changed, 210 insertions(+), 34 deletions(-) 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/CMakeLists.txt b/CMakeLists.txt index f87a439..ef483c8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -206,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 @@ -251,11 +260,9 @@ 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), and no -# reassociation either: oneAPI's icpx compiles the host with -fp-model=fast by -# default, which reorders the reference's sums and products. -set_source_files_properties(src/foldquant.cpp src/foldquant_ref.cpp PROPERTIES COMPILE_OPTIONS - "$<$:-fp-model=precise>;$<$:-Wno-overriding-option>;-ffp-contract=off") +# 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}) @@ -357,7 +364,7 @@ if(GGML_SYCL AND NOT GGML_CUDA) # rides in the executable's device image. 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;-fp-model=precise;-Wno-overriding-option;-ffp-contract=off") + COMPILE_OPTIONS "-fsycl;-ffp-contract=off") target_link_options(vla_core PUBLIC -fsycl "SHELL:-Xsycl-target-backend=spir64 \"-cl-fp32-correctly-rounded-divide-sqrt\"") target_compile_definitions(vla_core PRIVATE VLA_FQ_SYCL) diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md index 161f3d0..118494b 100644 --- a/docs/QUANTIZATION.md +++ b/docs/QUANTIZATION.md @@ -203,7 +203,10 @@ reference as long as the float steps are: the source is built with 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. A whole π0.5 model follows the CUDA integer +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. A whole π0.5 model follows the CUDA integer path to 1.00000 action cosine at W8A8, 0.9997 at W4A4 with INT8 o/down and 0.9991 at W4A4. 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 diff --git a/src/sycl/vla_sycl_foldquant.cpp b/src/sycl/vla_sycl_foldquant.cpp index a4b7926..4e1fc51 100644 --- a/src/sycl/vla_sycl_foldquant.cpp +++ b/src/sycl/vla_sycl_foldquant.cpp @@ -639,7 +639,23 @@ void launch_epilogue(sycl::queue & q, const GemmParams & g, const int32_t * acc, }); } +bool gemm_dnnl_impl(sycl::queue & q, const GemmParams & g, int wbits, int abits, bool has_bias, bool has_res); + +// oneDNN, unless it cannot drive this device (no engine, no int8 matmul): then +// this file's kernels, for good on that queue. 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[&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); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3869632..d8be335 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -83,9 +83,9 @@ 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: precise floating point, no contraction. - target_compile_options(test_foldquant_sycl_op PRIVATE -Wall -Wextra -fp-model=precise -Wno-overriding-option - -ffp-contract=off) + # 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() diff --git a/tests/test_foldquant_sycl_op.cpp b/tests/test_foldquant_sycl_op.cpp index 51ef686..6a51ec5 100644 --- a/tests/test_foldquant_sycl_op.cpp +++ b/tests/test_foldquant_sycl_op.cpp @@ -17,6 +17,12 @@ // 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" @@ -28,13 +34,22 @@ #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 { @@ -55,7 +70,14 @@ struct Case { // The shape under test; main() runs the cases over a list of them. int64_t K = 256, N = 128, T = 37; -int run_case(ggml_backend_t backend, const Case & c) { +// 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; @@ -93,13 +115,6 @@ int run_case(ggml_backend_t backend, const Case & c) { 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_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; @@ -119,19 +134,54 @@ int run_case(ggml_backend_t backend, const Case & c) { ggml_set_output(y); ggml_cgraph * gf = ggml_new_graph(C); ggml_build_forward_expand(gf, y); - ggml_gallocr_t 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; } - ggml_backend_tensor_set(x, hx.data(), 0, ggml_nbytes(x)); - if (r) ggml_backend_tensor_set(r, hr.data(), 0, ggml_nbytes(r)); - if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { - std::printf("FAIL: %s: graph compute\n", c.name); + 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); - 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)); + 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; @@ -166,9 +216,10 @@ int run_case(ggml_backend_t backend, const Case & c) { "%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)); - ggml_gallocr_free(alloc); + for (void * p : usm) sycl::free(p, *ex.q); + if (alloc) ggml_gallocr_free(alloc); ggml_free(C); - ggml_backend_buffer_free(wbuf); + if (wbuf) ggml_backend_buffer_free(wbuf); ggml_free(W); return ok ? 0 : 1; } @@ -176,10 +227,24 @@ int run_case(ggml_backend_t backend, const Case & c) { } // namespace int main() { - ggml_backend_t backend = ggml_backend_sycl_init(0); - if (!backend) { - std::printf("SKIP: no SYCL device\n"); - return 0; + 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[] = { @@ -212,10 +277,10 @@ int main() { 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(backend, c); + fails += run_case(ex, c); } } std::printf("%s\n", fails ? "test_foldquant_sycl_op: FAILED" : "test_foldquant_sycl_op: PASS"); - ggml_backend_free(backend); + if (ex.backend) ggml_backend_free(ex.backend); return fails ? 1 : 0; } From b32ecdaad1818dcff0f6d2efd5d2cd838a8aca2e Mon Sep 17 00:00:00 2001 From: hungho77 Date: Sat, 3 Oct 2026 11:54:36 +0700 Subject: [PATCH 08/11] docs: FoldQuant W8A8 numbers from a calibrated FoldQuantVLA checkpoint The pi0.5 W8A8 rows came from scripts/foldquant_fake_export.py (rotation and per-row RTN, no calibration). Re-measured with FoldQuantVLA's W8A8 arm (w8a8_sr + w8a8_sh) calibrated on the same 128 LIBERO frames as vrfai/pi05-libero-w4a4, converted byte-identical to its TensorRT graphs: SYCL 305 ms (bf16 326), OpenVINO GPU 655 ms / 4.7 GB, CPU 3653 ms / 8.6 GB; actions vs the CUDA integer path 0.99999 (SYCL), 1.00000 (OpenVINO). The tables now name the checkpoints they were measured on. --- docs/QUANTIZATION.md | 4 ++-- docs/backend/ov.md | 11 +++++++++-- docs/backend/sycl.md | 11 +++++++++-- 3 files changed, 20 insertions(+), 6 deletions(-) diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md index 118494b..5027b96 100644 --- a/docs/QUANTIZATION.md +++ b/docs/QUANTIZATION.md @@ -207,8 +207,8 @@ 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. A whole π0.5 model follows the CUDA integer -path to 1.00000 action cosine at W8A8, 0.9997 at W4A4 with INT8 o/down and 0.9991 -at W4A4. Both GPUs are exact to the same reference, so the gap comes from the +path to 0.99999 action cosine at W8A8, 0.9997 at W4A4 with INT8 o/down and 0.9991 +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. diff --git a/docs/backend/ov.md b/docs/backend/ov.md index bb38f97..e702ffc 100644 --- a/docs/backend/ov.md +++ b/docs/backend/ov.md @@ -304,11 +304,11 @@ B390 iGPU, OpenVINO 2026.4, `vla-bench` p50 over 10 calls after 3 warmups: | 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 (uncalibrated) | 3763 | 9.4 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 (uncalibrated) | 595 | 9.6 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 @@ -316,6 +316,13 @@ 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 d97b7d6..65b97ec 100644 --- a/docs/backend/sycl.md +++ b/docs/backend/sycl.md @@ -145,10 +145,17 @@ B390 iGPU, oneAPI 2026.1, `vla-bench` p50 over 10 calls after 3 warmups: | 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 (uncalibrated) | 4.73 GB | 330 | +| 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 runs at bf16 speed. +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 From 870786b5ad384b35acb4082a7dea13a5ab3a5a06 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Mon, 5 Oct 2026 10:19:39 +0700 Subject: [PATCH 09/11] CI fixes for the SYCL and OpenVINO jobs - FoldQuant SYCL: oneDNN only on GPU devices. On the CI runner's OpenCL CPU device (an AMD EPYC, no VNNI) oneDNN's int8 matmul sums pairs in saturating 16-bit lanes and the outputs were off; the native kernels (and oneDNN with s4 weights) were exact there. vla.cpp runs ggml-sycl on GPUs only. - test_graph_names: an OpenVINO build already defines GGML_USE_OPENVINO; redefining it broke the -Werror build. --- docs/QUANTIZATION.md | 5 +++-- src/sycl/vla_sycl_foldquant.cpp | 8 ++++++-- tests/test_graph_names.cpp | 2 ++ 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md index 5027b96..f140874 100644 --- a/docs/QUANTIZATION.md +++ b/docs/QUANTIZATION.md @@ -206,8 +206,9 @@ 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. 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.9991 +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.9991 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 diff --git a/src/sycl/vla_sycl_foldquant.cpp b/src/sycl/vla_sycl_foldquant.cpp index 4e1fc51..aafe04d 100644 --- a/src/sycl/vla_sycl_foldquant.cpp +++ b/src/sycl/vla_sycl_foldquant.cpp @@ -641,10 +641,14 @@ void launch_epilogue(sycl::queue & q, const GemmParams & g, const int32_t * acc, bool gemm_dnnl_impl(sycl::queue & q, const GemmParams & g, int wbits, int abits, bool has_bias, bool has_res); -// oneDNN, unless it cannot drive this device (no engine, no int8 matmul): then -// this file's kernels, for good on that queue. +// 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; 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 From bea469801b5e64c9c69d6cbc011b300a3731ebfc Mon Sep 17 00:00:00 2001 From: hungho77 Date: Mon, 5 Oct 2026 10:42:47 +0700 Subject: [PATCH 10/11] FoldQuant SYCL: one device image per kernel The CI runner's OpenCL CPU compiler crashed (free(): invalid pointer in sycl-kernel-reduce-cross-barrier-values) JIT-compiling the module, which holds every kernel, the XMX one included. With per-kernel images a device compiles only the kernels it launches. No change on the Arc GPU (W4A4 p50 291 ms); 3/3 CPU-device runs pass on the Intel PC. QUANTIZATION.md: W4A4 SYCL vs CUDA is 0.9995 now that host code is built with precise floating point under icpx (bf16 moved by 6e-4 with it). --- CMakeLists.txt | 6 ++++-- docs/QUANTIZATION.md | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index ef483c8..2a01c88 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -361,11 +361,13 @@ if(GGML_SYCL AND NOT GGML_CUDA) # 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. + # 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 + 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 diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md index f140874..033b1c5 100644 --- a/docs/QUANTIZATION.md +++ b/docs/QUANTIZATION.md @@ -208,7 +208,7 @@ 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.9991 +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 From 68c1f811590f5f9a3d179dbe2ae3df587b4431af Mon Sep 17 00:00:00 2001 From: hungho77 Date: Mon, 5 Oct 2026 13:01:27 +0700 Subject: [PATCH 11/11] FoldQuant SYCL: VLA_FQ_CHECK re-checks every node in a model run The CUDA path's diagnostic, ported: after each node, recompute it with the CPU reference on host copies of the same inputs and report byte differences. On an Arc B390, GR00T N1.5 W4A4 (732 nodes) and N1.6 W4A4 (1376 nodes) are byte-identical node for node. --- docs/QUANTIZATION.md | 4 +-- src/sycl/vla_sycl_foldquant.cpp | 52 +++++++++++++++++++++++++++++++-- 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md index 033b1c5..027f7d3 100644 --- a/docs/QUANTIZATION.md +++ b/docs/QUANTIZATION.md @@ -238,8 +238,8 @@ 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` 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 +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 diff --git a/src/sycl/vla_sycl_foldquant.cpp b/src/sycl/vla_sycl_foldquant.cpp index aafe04d..d25dbb0 100644 --- a/src/sycl/vla_sycl_foldquant.cpp +++ b/src/sycl/vla_sycl_foldquant.cpp @@ -37,6 +37,7 @@ #include "foldquant.h" #include "foldquant_ref.h" +#include "env_flag.h" #include "ggml.h" #include @@ -771,15 +772,60 @@ 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); - if (m == FQ_ACT_MAGIC) - return run_act(q, dst, *static_cast(ud)); + 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); - return run_gemm(q, dst, s, blob_abits(dst->src[1], s.K)); + 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) {