diff --git a/CHANGELOG.md b/CHANGELOG.md index 42ea5ff0..91d844b4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,8 +35,13 @@ Notable changes to vla.cpp. Format loosely follows [Keep a Changelog](https://ke reference (1.0e-3 served): pads that read a weight, also answered `NO_SUPPORT`, run on the CPU, and its ResNet's im2col matrices are zero-padded to a multiple of 32 so the convs' matmuls take the HMX (152.9 ms, - 170.8 ms on the CPU). GR00T N1.5 and N1.6 are above the bar on a served LIBERO - observation, as the CPU is at F16. See + 170.8 ms on the CPU). Explicit attention takes F16 K/V with its head dim and + keys zero-padded to a multiple of 32, the padding masked out, so both matmuls + run on the HMX: OpenVLA-OFT's LM, Octo's block transformer, Evo-1's DiT + cross-attention and GR00T's DiT (OpenVLA-OFT 35% faster, Evo-1 29%, Octo + 29%, GR00T N1.5 11%); VLA-Adapter's LM only for V. Fidelity is measured + against a CPU-only build of the same commit. GR00T N1.5 and N1.6 are above + the bar on a served LIBERO observation, as the CPU is at F16. See `docs/backend/hexagon.md` for the build, the skel install and the numbers. - **ROCm/HIP backend for VLA inference.** `-DGGML_HIP=ON` selects ggml's HIP backend without compiling vla.cpp's CUDA-only kernels. On Linux `gfx1200` diff --git a/README.md b/README.md index a6197876..7b3e6873 100644 --- a/README.md +++ b/README.md @@ -105,12 +105,6 @@ Use [ggml-compiler](https://github.com/khanhnd61-vr/ggml-compiler) and follow [t Legend: `Y` = supported (released and benchmarked), `~` = in progress, `-` = planned. -The Hexagon column is the Dragonwing IQ-10 under Linux -([docs/backend/hexagon.md](docs/backend/hexagon.md)). On a Snapdragon X laptop -under Windows, π0 and VLA-JEPA land above the 2.9e-3 fidelity bar, OpenVLA-OFT -was not attempted (about 14 GB at F16, over its 8 GB budget) and ACT was not run -([docs/backend/hexagon-windows.md](docs/backend/hexagon-windows.md)). - --- ## Rollout on a real robot diff --git a/docs/backend/hexagon-windows.md b/docs/backend/hexagon-windows.md index 0241be4c..c259e8d2 100644 --- a/docs/backend/hexagon-windows.md +++ b/docs/backend/hexagon-windows.md @@ -130,7 +130,7 @@ PyTorch and LIBERO; the smoke test used a 20-line pyzmq client instead. | Variable / flag | Effect in the Hexagon and OpenCL builds | |---|---| -| `VLA_DEVICE=cpu` | Skip the accelerator. The same binary then produces the CPU reference. | +| `VLA_DEVICE=cpu` | Skip the accelerator. Not a fidelity reference: the Hexagon build's NPU changes (F16 K/V, padding) run on the CPU too, so compare against a CPU-only build with `--weight-dtype f32 --flash-attn 0`. | | `VLA_DEVICE=` | Registry device index: an HTP session, or an OpenCL GPU | | `--weight-dtype f16\|f32\|bf16` | Resident dtype for GEMM weights. New: `f16`, which is also the default on both accelerators. | | `--flash-attn on\|off` | On by default in Hexagon builds, off elsewhere | diff --git a/docs/backend/hexagon.md b/docs/backend/hexagon.md index 14a4abd8..d87a1ee2 100644 --- a/docs/backend/hexagon.md +++ b/docs/backend/hexagon.md @@ -143,11 +143,11 @@ the same path. Flash attention on the HMX (v81) keeps scores and probabilities in F16. That is a precision limit rather than a wrong op, but in GR00T N1.7's LM it moves the -action chunk 5.4e-3 from the CPU reference. N1.7's LM therefore takes the +action chunk 6.2e-3 from the CPU reference. N1.7's LM therefore takes the explicit path unless `--flash-attn` is given: 2.0e-3, at the same latency. VLA-JEPA's LM does the same: with flash attention a served LIBERO observation -lands 6.7e-3 from the reference (the CPU with flash attention 5.5e-3), without -it 2.6e-3, at a cost of 683 to 810 ms with two views. +lands 6.2e-3 from the reference, without it 2.2e-3, at a cost of 683 to 810 ms +with two views. The softmax takes rows of up to 32 elements or a multiple of 32. Where a sequence can grow without changing its result, vla.cpp pads it to @@ -165,7 +165,7 @@ tower takes two camera views that way: its action chunk lands 0.16 from the CPU reference. The wrapper runs those matmuls on the CPU. On Hexagon `linear()` in `src/layers/linear.h` multiplies a contiguous batch as one 2D matrix of rows instead, the same product, so most of them never reach the rule: GR00T N1.6 -with two views lands 8.8e-4 from the reference, and the matmuls of TurboVLA's +with two views lands 7.5e-4 from the reference, and the matmuls of TurboVLA's DINOv3, which takes its two views as a batch, run on the NPU. A copy that transposes the innermost dim, such as the `CONT` of a permuted V @@ -196,6 +196,16 @@ Hexagon builds: 311 ms. - **Evo-1's LM attention multiplied F32 K and V against 14 query heads.** K and V now go to F16 (`fa_kv()`); the scores still accumulate in F32. +- **Explicit attention elsewhere multiplied F32 K and V too**, and its key + counts and head dims are often not multiples of 32. OpenVLA-OFT's LM (330 + keys), Octo's block transformer, Evo-1's DiT cross-attention (head dim 112, + 1025 keys) and GR00T's DiT (head dim 48) now take F16 K/V with the head dim + and the keys zero-padded to `gemm_width()` and the padded keys masked out + (`hmx_kv()` in `dit_head.cpp`): OpenVLA-OFT 1725 to 1120 ms, Evo-1 2348 to + 1661 ms, Octo-Small 119 to 85 ms, GR00T N1.5 595 to 532 ms. Two LMs keep part + of it in F32: in VLA-Adapter's LM only V goes to F16, since F16 scores put the + chunk 3.9e-3 off, and VLA-JEPA's cross-attention K/V overflow F16 (NaN + actions), so only its DiT self-attention takes the HMX. - **ACT's ResNet convs are im2col and one matmul** with the patches as its first operand. The stem's K of 147 and layer3's and layer4's 1200 and 300 output positions (640x480) are not multiples of 32: 80 ms of the 101 ms spent @@ -206,10 +216,10 @@ Hexagon builds: benchmark's protocol ACT takes 152.9 ms on the NPU against 170.8 ms on the CPU. -What is left on the vector units: GR00T's DiT attention (head dim 48) and Evo-1's -DiT cross-attention (head dim 112, 1025 keys), whose shapes the HMX does not -take, and Evo-1's InternViT attention, which stays on the CPU because F16 K/V -there measured 3.9e-2. +What is left on the vector units: VLA-Adapter's QK scores and VLA-JEPA's LM and +DiT cross-attention scores, which stay F32 for the reasons above, and Evo-1's +InternViT attention, which stays on the CPU because F16 K/V there measured +3.9e-2. ## What changed in `vla.cpp` @@ -220,7 +230,8 @@ The shape of the diff was predicted correctly, apart from the scheduler: 2. **A rung in [src/backend.h](../../src/backend.h).** - It goes through the registry, because `ggml_backend_hexagon_init` is declared but not defined upstream. - - `VLA_DEVICE` indexes HTP sessions, and `VLA_DEVICE=cpu` gives the reference. + - `VLA_DEVICE` indexes HTP sessions, and `VLA_DEVICE=cpu` skips them. That + run is not a fidelity reference (see [Fidelity](#fidelity)). - It sets the defaults above and points `ADSP_LIBRARY_PATH` at the executable's folder. 3. **Weights.** @@ -254,9 +265,11 @@ found. to the CPU and how much was copied (`VLA_FALLBACK_STATS=1` for per-graph detail). 4. **Fidelity before latency.** - - Compare `vla_predict_check` output against the same binary with - `VLA_DEVICE=cpu`, and against `--weight-dtype f32`. Add `--flash-attn 0` - to the reference: a Hexagon build defaults to flash attention on the CPU too. + - Compare `vla_predict_check` output against a CPU-only build + (`-DGGML_HEXAGON=OFF`) of the same commit, run with `--weight-dtype f32 + --flash-attn 0`. The Hexagon binary with `VLA_DEVICE=cpu` is not a + reference: its NPU changes (flash attention, F16 K/V, padding to + `gemm_width()`) run on the CPU too. - The bar is 2.9e-3. - A run with a wrong kernel still prints a latency, and it is usually a good one. @@ -355,30 +368,31 @@ export LD_LIBRARY_PATH=$PWD/pkg/lib ### Fidelity -`vla_predict_check` at the Hexagon defaults, against the same binary with -`VLA_DEVICE=cpu --weight-dtype f32 --flash-attn 0`. A Hexagon build turns flash -attention on by default even on the CPU, so the reference asks for the explicit -path that other builds run. Views are the ones the LIBERO client sends, so they +`vla_predict_check` at the Hexagon defaults, against a CPU-only build +(`-DGGML_HEXAGON=OFF`) of the same commit with `--weight-dtype f32 +--flash-attn 0`. The Hexagon build's own CPU run is not a reference: what it +changes for the NPU (F16 K/V, padding to `gemm_width()`) is compiled in and runs +on the CPU too. Views are the ones the LIBERO client sends, so they differ from the benchmark's for VLA-JEPA, VLA-Adapter, OpenVLA-OFT and Octo: | Model | Views | max\|Δ\| | |---|--:|--:| -| SmolVLA | 2 | 7.2e-4 | -| π0 | 2 | 1.2e-3 | +| SmolVLA | 2 | 7.4e-4 | +| π0 | 2 | 7.7e-4 | | π0.5 | 2 | 4.8e-4 | -| Evo-1 | 1 | 5.5e-4 | -| GR00T N1.5 | 1 | 2.4e-3 | -| GR00T N1.6 | 1 | 1.4e-3 | +| Evo-1 | 1 | 6.2e-4 | +| GR00T N1.5 | 1 | 1.9e-3 | +| GR00T N1.6 | 1 | 1.3e-3 | | GR00T N1.7 | 1 | 2.0e-3 | -| VLA-JEPA | 2 | 2.2e-3 | -| VLA-Adapter | 2 | 5.4e-4 | -| OpenVLA-OFT | 2 | 3.2e-4 | -| Octo-Small | 1 | 5.3e-4 | +| VLA-JEPA | 2 | 2.1e-3 | +| VLA-Adapter | 2 | 3.9e-4 | +| OpenVLA-OFT | 2 | 3.1e-4 | +| Octo-Small | 1 | 5.0e-4 | | TurboVLA | 2 | 3.7e-4 | | ACT | 1 | 6.1e-4 | -GR00T N1.6 with two views is 8.8e-4. At the benchmark's single view VLA-JEPA is -1.9e-3, VLA-Adapter 5.7e-4 and OpenVLA-OFT 2.1e-4. ACT takes its one 640x480 +GR00T N1.6 with two views is 7.5e-4. At the benchmark's single view VLA-JEPA is +1.8e-3, VLA-Adapter 4.8e-4 and OpenVLA-OFT 1.8e-4. ACT takes its one 640x480 view, the benchmark's input. ### Serving @@ -387,7 +401,8 @@ view, the benchmark's input. `eval/client/vla_cpp_client.py` and the adapters `run_sim_client_direct.py` uses, on one fixed LIBERO observation (two views, the task prompt, a fixed robot state), with fixed noise for the archs that draw it. The same request goes to a -Hexagon server and to a CPU server with `--weight-dtype f32 --flash-attn 0`. +Hexagon server and to a server from the CPU-only build with `--weight-dtype f32 +--flash-attn 0`. Every model returns a finite chunk, the same chunk for a repeated request, an error reply to a malformed or an invalid request, and keeps serving afterwards. ACT has no `eval/client` preset; it gets the same checks from raw `vla.proto` @@ -395,23 +410,23 @@ requests carrying one fixed 640x480 view and a fixed 14-dim state. | Model | max\|Δ\| vs CPU F32 | CPU F16 vs CPU F32 | |---|--:|--:| -| SmolVLA | 1.3e-3 | | -| π0.5 | 3.5e-4 | | +| SmolVLA | 1.9e-3 | | +| π0.5 | 3.7e-4 | | | Evo-1 | 1.1e-3 | | -| GR00T N1.5 | 1.1e-2 | 4.3e-2 | -| GR00T N1.6 | 1.3e-2 | 1.1e-2 | -| GR00T N1.7 | 2.5e-3 | | +| GR00T N1.5 | 1.2e-2 | 2.8e-2 | +| GR00T N1.6 | 1.0e-2 | 9.0e-3 | +| GR00T N1.7 | 2.0e-3 | | | π0 | 5.8e-4 | | -| VLA-JEPA | 2.6e-3 | 2.1e-2 | -| VLA-Adapter | 5.4e-4 | | -| OpenVLA-OFT | 5.9e-4 | | -| Octo-Small | 2.4e-4 | | +| VLA-JEPA | 2.2e-3 | | +| VLA-Adapter | 3.9e-4 | | +| OpenVLA-OFT | 7.5e-4 | | +| Octo-Small | 1.4e-4 | | | TurboVLA | 8.5e-4 | | | ACT | 1.0e-3 | | GR00T N1.5 and N1.6 are above the bar on this input, and so is the CPU itself with F16 weights: N1.5 lands closer to the reference than the CPU at F16, N1.6 -slightly further (1.3e-2 against 1.1e-2). They are sensitive to 16-bit precision +slightly further (1.0e-2 against 9.0e-3). They are sensitive to 16-bit precision rather than miscomputed. Running them with F32 weights on the NPU has not been tried yet. diff --git a/docs/benchmark/README.md b/docs/benchmark/README.md index 50c184ea..50791593 100644 --- a/docs/benchmark/README.md +++ b/docs/benchmark/README.md @@ -12,7 +12,7 @@ AGX Orin row ran at `e6653e8`, on the bit-exact CUDA path ran at `d0e5d00`, PR #35 merged with main `f464653` (same llama.cpp pin), with fixed-input checks against the CPU build in [backend/rocm.md](../backend/rocm.md). The Qualcomm Dragonwing IQ-10 row ran at -`7c76ea9`, its ACT cell at `0f3e010` (same llama.cpp pin). Each report +`bd166b8` (same llama.cpp pin). Each report lists the device, the build, each model's configuration, latency and peak memory at the defaults, and the fastest runtime flags for that device. @@ -36,7 +36,7 @@ means the model does not fit or fails there; each report says which and why. | [Intel Arc B390 iGPU (OpenVINO)](core-ultra-x7-358h.md) | N/A | n/r | 28.5 | 116.3 | 170.2 | 106.0 | 122.8 | 131.4 | N/A | 213.1 | 620.5 | 232.5 | 1262.5 | 566.0 | | [Intel Arc A380 (SYCL)](arc-a380.md) | 60.5 | n/r | 103.5 | 355.9 | 408.9 | 472.0 | 646.0 | 615.3 | N/A | 734.5 | 921.9 | 945.6 | - | 1047.3 | | [Snapdragon X Hexagon NPU](snapdragon-x-hexagon.md) | 239.7 | n/r | 2280.8 | 821.0 | 4072.1 | 2755.7 | 2940.0 | 1400.2 | N/A | 1239.3 | 5730.0 | 9321.7 | - | 7208.6 | -| [Qualcomm Dragonwing IQ-10 Hexagon NPU](qualcomm-iq10.md) | 117.6 | 148.5 | 188.1 | 505.4 | 703.2 | 591.6 | 696.5 | 759.5 | N/A | 583.6 | 1114.0 | 1113.0 | 1720.3 | 2333.6 | +| [Qualcomm Dragonwing IQ-10 Hexagon NPU](qualcomm-iq10.md) | 81.7 | 147.8 | 184.6 | 489.7 | 664.0 | 530.2 | 680.9 | 739.6 | N/A | 583.3 | 1117.6 | 1114.0 | 1117.7 | 1653.5 | | [Intel AI Boost NPU (OpenVINO)](core-ultra-x7-358h.md) | N/A | n/r | 23.5 | 90.7 | 126.7 | 134.0 | 161.1 | 144.9 | N/A | 161.3 | 438.3 | 430.2 | 405.8 | 256.3 | | [Intel Core i7-14700F (CPU)](core-i7-14700f.md) | 39.8 | n/r | 195.1 | 825.5 | 1140.6 | 1913.8 | 1627.9 | 1453.4 | N/A | 1689.2 | 6455.7 | 6068.1 | 9869.0 | 4159.6 | | [Intel Core i9-14900HX (CPU)](core-i9-14900hx.md) | 67.0 | n/r | 314.7 | 1227.2 | 1659.3 | 1935.4 | 1688.9 | 1505.5 | N/A | 1822.6 | 6182.1 | 5735.0 | 9126.4 | 4118.2 | diff --git a/docs/benchmark/qualcomm-iq10.md b/docs/benchmark/qualcomm-iq10.md index 39d1ef14..d583607b 100644 --- a/docs/benchmark/qualcomm-iq10.md +++ b/docs/benchmark/qualcomm-iq10.md @@ -15,9 +15,9 @@ through vla.cpp's fallback wrapper. | Power | CPU governor `performance` | | OS | Ubuntu 24.04.4 LTS, kernel 6.6.110 | | Toolchain | `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7`: clang 21.1.8, Hexagon SDK 6.6.0.0 (tools 19.0.07) | -| Commit | `7c76ea9`; ACT at `0f3e010` | +| Commit | `bd166b8` | | llama.cpp | `b11223` | -| Date | 2026-11-05 14:26 to 14:57 (UTC+07); ACT 17:00 to 17:01 | +| Date | 2026-11-06 10:46 to 11:16 (UTC+07) | | Build | `scripts/build_linux_hexagon.sh` (servers included); the HTP skel processed for the board and installed in `/vendor/dsp/cdsp0..3` | | Runtime flags | defaults; the fastest flags per model are in the second table | | Method | 3 warmups + 20 timed reps per process, 3 processes. The process with the lowest mean is reported, and memory is the peak over all three. | @@ -69,19 +69,19 @@ vision stage separately. | Model | Views | Input | min ms | mean ms | p50 ms | p90 ms | vision ms | Peak RSS MiB | System RAM MiB | |---|--:|--:|--:|--:|--:|--:|--:|--:|--:| -| Octo-Small | 2 | 256 + 128 | 117.6 | 119.4 | 118.4 | 120.7 | 21.2 | 253 | 723 | -| ACT | 1 | 640 x 480 | 148.5 | 152.9 | 152.5 | 157.3 | - | 194 | 467 | -| TurboVLA | 2 | 256 | 188.1 | 190.8 | 188.6 | 194.9 | - | 249 | 1075 | -| VLA-JEPA | 1 | 256 | 505.4 | 512.6 | 510.1 | 522.2 | 196.4 | 184 | 3969 | -| SmolVLA | 2 | 512 | 583.6 | 584.4 | 584.4 | 585.1 | 229.9 | 145 | 1196 | -| GR00T N1.5 | 1 | 224 | 591.6 | 594.9 | 593.4 | 597.1 | 138.9 | 162 | 3639 | -| GR00T N1.6 | 1 | 224 | 696.5 | 699.3 | 699.2 | 702.2 | 145.1 | 147 | 4692 | -| VLA-Adapter | 1 | 224 | 703.2 | 708.4 | 708.2 | 715.3 | 258.1 | 872 | 3566 | -| GR00T N1.7 | 1 | 256 | 759.5 | 766.3 | 766.0 | 773.4 | 197.6 | 191 | 5113 | -| π0 | 2 | 224 | 1114.0 | 1116.6 | 1116.4 | 1118.2 | 277.1 | 284 | 5537 | -| π0.5 | 2 | 224 | 1113.0 | 1117.3 | 1116.6 | 1120.7 | 277.9 | 285 | 5826 | -| OpenVLA-OFT | 1 | 224 | 1720.3 | 1725.5 | 1724.7 | 1731.0 | 263.4 | 844 | 15451 | -| Evo-1 | 1 | 448 | 2333.6 | 2348.3 | 2349.9 | 2360.1 | 820.6 | 319 | 1757 | +| Octo-Small | 2 | 256 + 128 | 81.7 | 84.6 | 82.3 | 89.6 | 21.9 | 253 | 728 | +| ACT | 1 | 640 x 480 | 147.8 | 151.2 | 150.7 | 153.6 | - | 193 | 471 | +| TurboVLA | 2 | 256 | 184.6 | 191.8 | 190.3 | 199.4 | - | 249 | 1118 | +| VLA-JEPA | 1 | 256 | 489.7 | 495.8 | 494.7 | 501.6 | 196.4 | 184 | 4023 | +| GR00T N1.5 | 1 | 224 | 530.2 | 532.0 | 531.3 | 534.3 | 138.5 | 162 | 3734 | +| SmolVLA | 2 | 512 | 583.3 | 584.0 | 583.9 | 584.7 | 229.8 | 146 | 1194 | +| VLA-Adapter | 1 | 224 | 664.0 | 668.6 | 666.6 | 673.7 | 257.8 | 872 | 3584 | +| GR00T N1.6 | 1 | 224 | 680.9 | 682.1 | 681.6 | 683.9 | 145.3 | 146 | 4889 | +| GR00T N1.7 | 1 | 256 | 739.6 | 746.4 | 745.1 | 753.8 | 196.4 | 190 | 5215 | +| π0.5 | 2 | 224 | 1114.0 | 1116.1 | 1115.5 | 1117.8 | 278.0 | 285 | 5802 | +| π0 | 2 | 224 | 1117.6 | 1120.5 | 1119.9 | 1123.8 | 277.8 | 284 | 5547 | +| OpenVLA-OFT | 1 | 224 | 1117.7 | 1120.2 | 1120.4 | 1121.8 | 262.3 | 844 | 15321 | +| Evo-1 | 1 | 448 | 1653.5 | 1661.4 | 1661.2 | 1665.6 | 815.6 | 319 | 1556 | ### Fastest configuration @@ -94,7 +94,7 @@ under the full protocol: The weight dtype stays F16 in every set. `--weight-dtype f16` is the default, and `--weight-dtype bf16` is left out: the NPU has no BF16 kernels, so every -matmul would run on the CPU (SmolVLA: 18.7 s against 0.82 s). +matmul would run on the CPU (SmolVLA: 18.6 s against 0.58 s, mean, 5 reps). The screen runs 2 warmups and 5 reps in one process per flag set. A flag set replaces the defaults only if its mean is more than 3% lower, both in the screen @@ -103,42 +103,46 @@ and in the full re-run. Flash attention and the matmul precision can move the action chunk. Success rates measured at the defaults therefore do not carry over to these rows. For VLA-JEPA and GR00T N1.7, `--flash-attn` puts their LM back on flash attention, -which the defaults leave off because it is outside the fidelity bar (VLA-JEPA -6.7e-3 on a served LIBERO observation, N1.7 5.4e-3 on fixed inputs). +which the defaults leave off because it is outside the fidelity bar against a +CPU-only build: VLA-JEPA 6.2e-3 on a served LIBERO observation (2.7e-3 on fixed +inputs), N1.7 6.2e-3 on fixed inputs (2.7e-3 served). | Model | Fastest flags | min ms | mean ms | p50 ms | p90 ms | vision ms | Peak RSS MiB | System RAM MiB | vs defaults | |---|---|--:|--:|--:|--:|--:|--:|--:|--:| -| Octo-Small | *(defaults)* | 117.6 | 119.4 | 118.4 | 120.7 | 21.2 | 253 | 723 | - | -| ACT | *(defaults)* | 148.5 | 152.9 | 152.5 | 157.3 | - | 194 | 467 | - | -| TurboVLA | *(defaults)* | 188.1 | 190.8 | 188.6 | 194.9 | - | 249 | 1075 | - | -| VLA-JEPA | `--flash-attn --mm-prec default` | 442.1 | 452.4 | 450.2 | 460.6 | 198.3 | 184 | 3985 | -11.7% | -| SmolVLA | *(defaults)* | 583.6 | 584.4 | 584.4 | 585.1 | 229.9 | 145 | 1196 | - | -| GR00T N1.5 | *(defaults)* | 591.6 | 594.9 | 593.4 | 597.1 | 138.9 | 162 | 3639 | - | -| GR00T N1.6 | *(defaults)* | 696.5 | 699.3 | 699.2 | 702.2 | 145.1 | 147 | 4692 | - | -| VLA-Adapter | *(defaults)* | 703.2 | 708.4 | 708.2 | 715.3 | 258.1 | 872 | 3566 | - | -| GR00T N1.7 | `--flash-attn --mm-prec default` | 727.8 | 734.1 | 733.2 | 742.4 | 196.2 | 191 | 5081 | -4.2% | -| π0 | *(defaults)* | 1114.0 | 1116.6 | 1116.4 | 1118.2 | 277.1 | 284 | 5537 | - | -| π0.5 | *(defaults)* | 1113.0 | 1117.3 | 1116.6 | 1120.7 | 277.9 | 285 | 5826 | - | -| OpenVLA-OFT | *(defaults)* | 1720.3 | 1725.5 | 1724.7 | 1731.0 | 263.4 | 844 | 15451 | - | -| Evo-1 | *(defaults)* | 2333.6 | 2348.3 | 2349.9 | 2360.1 | 820.6 | 319 | 1757 | - | +| Octo-Small | *(defaults)* | 81.7 | 84.6 | 82.3 | 89.6 | 21.9 | 253 | 728 | - | +| ACT | *(defaults)* | 147.8 | 151.2 | 150.7 | 153.6 | - | 193 | 471 | - | +| TurboVLA | *(defaults)* | 184.6 | 191.8 | 190.3 | 199.4 | - | 249 | 1118 | - | +| VLA-JEPA | `--flash-attn` | 427.8 | 434.2 | 432.6 | 441.7 | 196.1 | 184 | 3971 | -12.4% | +| GR00T N1.5 | *(defaults)* | 530.2 | 532.0 | 531.3 | 534.3 | 138.5 | 162 | 3734 | - | +| SmolVLA | *(defaults)* | 583.3 | 584.0 | 583.9 | 584.7 | 229.8 | 146 | 1194 | - | +| VLA-Adapter | *(defaults)* | 664.0 | 668.6 | 666.6 | 673.7 | 257.8 | 872 | 3584 | - | +| GR00T N1.6 | *(defaults)* | 680.9 | 682.1 | 681.6 | 683.9 | 145.3 | 146 | 4889 | - | +| GR00T N1.7 | `--flash-attn --mm-prec default` | 715.1 | 719.2 | 718.7 | 722.9 | 195.5 | 190 | 5074 | -3.6% | +| π0.5 | *(defaults)* | 1114.0 | 1116.1 | 1115.5 | 1117.8 | 278.0 | 285 | 5802 | - | +| π0 | *(defaults)* | 1117.6 | 1120.5 | 1119.9 | 1123.8 | 277.8 | 284 | 5547 | - | +| OpenVLA-OFT | *(defaults)* | 1117.7 | 1120.2 | 1120.4 | 1121.8 | 262.3 | 844 | 15321 | - | +| Evo-1 | *(defaults)* | 1653.5 | 1661.4 | 1661.2 | 1665.6 | 815.6 | 319 | 1556 | - |
Screen means (ms) per flag set | Model | defaults | `--flash-attn` | `--mm-prec default` | `--flash-attn --mm-prec default` | `--flash-attn 0` | |---|--:|--:|--:|--:|--:| -| Octo-Small | 121.8 | 120.4 | 121.8 | 123.2 | 123.8 | -| ACT | 143.4 | 153.4 | 157.5 | 155.8 | 174.8 | -| TurboVLA | 196.5 | 189.3 | 192.3 | 190.9 | 257.4 | -| VLA-JEPA | 514.8 | 454.3 | 512.0 | 452.9 | 1662.4 | -| SmolVLA | 584.6 | 585.0 | 585.2 | 584.6 | 2147.2 | -| GR00T N1.5 | 594.7 | 592.6 | 600.1 | 601.0 | 877.2 | -| GR00T N1.6 | 709.6 | 714.5 | 711.1 | 713.1 | 875.3 | -| VLA-Adapter | 712.4 | 711.1 | 713.7 | 712.8 | 989.2 | -| GR00T N1.7 | 765.5 | 739.0 | 758.3 | 733.1 | 865.3 | -| π0 | 1131.6 | 1138.1 | 1118.0 | 1128.8 | 2051.6 | -| π0.5 | 1125.8 | 1109.1 | 1134.2 | 1134.6 | 2036.0 | -| OpenVLA-OFT | 1737.9 | 1722.4 | 1715.7 | 1718.5 | 2005.8 | -| Evo-1 | 2345.2 | 2352.4 | 2355.4 | 2361.4 | 3877.1 | +| Octo-Small | 90.8* | 83.0 | 83.2 | 85.3 | 88.1 | +| ACT | 154.4 | 158.4 | 154.7 | 153.2 | 179.9 | +| TurboVLA | 192.0 | 192.6 | 194.0 | 191.9 | 257.3 | +| VLA-JEPA | 498.2 | 431.3 | 494.0 | 439.8 | 602.4 | +| GR00T N1.5 | 534.0 | 526.4 | 527.0 | 526.4 | 811.6 | +| SmolVLA | 584.4 | 584.7 | 585.2 | 585.4 | 2145.2 | +| VLA-Adapter | 671.6 | 670.4 | 669.4 | 671.9 | 950.8 | +| GR00T N1.6 | 689.2 | 685.8 | 687.3 | 683.0 | 851.0 | +| GR00T N1.7 | 758.3 | 733.1 | 752.8 | 723.2 | 858.2 | +| π0.5 | 1125.2 | 1117.6 | 1134.2 | 1134.2 | 2039.2 | +| π0 | 1113.2 | 1126.1 | 1114.8 | 1126.1 | 2056.3 | +| OpenVLA-OFT | 1125.5 | 1127.8 | 1123.4 | 1123.5 | 1402.2 | +| Evo-1 | 1668.8 | 1669.2 | 1669.8 | 1657.4 | 3116.8 | + +\* Octo-Small's defaults screen was re-run: the first run (672.8 ms, with a +241 ms vision stage) overlapped a transient slowdown on the board.
@@ -154,7 +158,8 @@ The HTP skel has to be processed for this board and installed in [backend/hexagon.md](../backend/hexagon.md#linux-iq-10). GR00T N1.6 runs one view here, as the configuration above fixes. With two views -one batched `MUL_MAT` of its SigLIP2 tower still runs on the CPU; see +its SigLIP2 tower's patch matmul and flash attention, batched over the views, +still run on the CPU; see [backend/hexagon.md](../backend/hexagon.md#op-coverage-measured). These rows time the defaults; they do not vouch for accuracy. On a served LIBERO diff --git a/src/models/evo1.cpp b/src/models/evo1.cpp index ce31ffab..547927ea 100644 --- a/src/models/evo1.cpp +++ b/src/models/evo1.cpp @@ -85,6 +85,7 @@ struct Evo1ModelArch : public ModelArchBase { struct MainIO { ggml_tensor *t_embeds=nullptr,*t_pos=nullptr,*t_lmmask=nullptr,*t_qmask=nullptr; ggml_tensor *t_state=nullptr,*t_x=nullptr,*t_amask=nullptr,*x_action=nullptr; + ggml_tensor *t_dmask=nullptr; std::vector lm_ok; }; graph_cache main_graph; @@ -870,6 +871,14 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { const int64_t E = embed_dim, hd_dit = E/dit_heads; const float scale_dit = 1.0f/std::sqrt((float) hd_dit); const int64_t Nctx = SEQ+1; + // On Hexagon the DiT's cross-attention pads its head dim and its keys with + // zeros to gemm_width(), the keys masked out, and takes F16 K/V: both + // matmuls then run on the HMX instead of as F32 x F32 on the vector units. + const int64_t HP = gemm_width(hd_dit), NK = gemm_width(Nctx); + ggml_tensor * t_dmask = nullptr; + if (NK > Nctx) { + t_dmask = ggml_new_tensor_2d(C, GGML_TYPE_F32, NK, horizon); ggml_set_input(t_dmask); ggml_set_output(t_dmask); + } ggml_tensor * t_embeds = ggml_new_tensor_2d(C, GGML_TYPE_F32, lm_hidden, SEQ); ggml_set_input(t_embeds); ggml_tensor * t_pos = ggml_new_tensor_1d(C, GGML_TYPE_I32, SEQ); ggml_set_input(t_pos); @@ -935,6 +944,12 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { } else { dc[i].K = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, kp, hd_dit, dit_heads, Nctx), 0, 2, 1, 3)); dc[i].V = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, vp, hd_dit, dit_heads, Nctx), 1, 2, 0, 3)); + if (HP > hd_dit || NK > Nctx) { + dc[i].K = ggml_pad(C, dc[i].K, (int) (HP-hd_dit), (int) (NK-Nctx), 0, 0); + dc[i].V = ggml_pad(C, dc[i].V, (int) (NK-Nctx), (int) (HP-hd_dit), 0, 0); + } + dc[i].K = fa_kv(C, dc[i].K); + dc[i].V = fa_kv(C, dc[i].V); } } @@ -1003,11 +1018,16 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { } else { ggml_tensor * qp = as_type(C, ggml_add(C, mm_act(C, c.Wq, x_q, at), c.bq), GGML_TYPE_F32); ggml_tensor * Q = ggml_cont(C, ggml_permute(C, ggml_reshape_3d(C, qp, hd_dit, dit_heads, horizon), 0, 2, 1, 3)); + if (HP > hd_dit) + Q = ggml_pad(C, Q, (int) (HP-hd_dit), 0, 0, 0); ggml_tensor * kq = ggml_mul_mat(C, c.K, Q); ggml_prec_set_acc(kq, GGML_PREC_F32); // The Evo-1 reference cross-attends over the full padded context - // (no key mask), so the action queries see every LM position. - ggml_tensor * aw = ggml_soft_max_ext(C, kq, nullptr, scale_dit, 0.0f); + // (no key mask), so the action queries see every LM position; + // t_dmask hides only the keys padded on Hexagon. + ggml_tensor * aw = ggml_soft_max_ext(C, kq, t_dmask, scale_dit, 0.0f); ggml_tensor * kqv = ggml_mul_mat(C, c.V, aw); + if (HP > hd_dit) + kqv = ggml_view_3d(C, kqv, hd_dit, horizon, dit_heads, kqv->nb[1], kqv->nb[2], 0); att_pre = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, kqv, 0, 2, 1, 3)), E, horizon); } ggml_tensor * attn_out = ggml_add(C, mm_act(C, w.Wo, as_type(C, att_pre, at), at), w.bo); @@ -1040,7 +1060,7 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { ggml_set_output(x_action); gio.t_embeds=t_embeds; gio.t_pos=t_pos; gio.t_lmmask=t_lmmask; gio.t_qmask=t_qmask; - gio.t_state=t_state; gio.t_x=t_x; gio.t_amask=t_amask; gio.x_action=x_action; + gio.t_state=t_state; gio.t_x=t_x; gio.t_amask=t_amask; gio.x_action=x_action; gio.t_dmask=t_dmask; ggml_cgraph * gf = ggml_new_graph_custom(C, main_nodes, false); ggml_build_forward_expand(gf, x_action); @@ -1067,6 +1087,14 @@ std::vector Evo1ModelArch::predict(const Inputs& in) { ggml_backend_tensor_set(t_amask, am.data(), 0, ggml_nbytes(t_amask)); } { std::vector qm(SEQ, 0.0f); for (int64_t p=0; pne[0], rows = gio.t_dmask->ne[1], Nctx = SEQ+1; + std::vector dm((size_t) (NK*rows), 0.0f); + for (int64_t r=0; r::infinity(); + ggml_backend_tensor_set(gio.t_dmask, dm.data(), 0, ggml_nbytes(gio.t_dmask)); + } gio.lm_ok = attn_ok; } ggml_backend_tensor_set(t_state, state_norm.data(), 0, ggml_nbytes(t_state)); diff --git a/src/models/octo.cpp b/src/models/octo.cpp index 4774c722..4e6cf65d 100644 --- a/src/models/octo.cpp +++ b/src/models/octo.cpp @@ -1149,8 +1149,11 @@ bool run_transformer_graph(OctoRuntime& rt, ggml_tensor * n1 = layer_norm(C, x, blk_w[i][0], blk_w[i][1], ln_eps); ggml_tensor * qkv = linear(C, blk_w[i][2], blk_w[i][3], n1); ggml_tensor * Q = ggml_cont(C, ggml_permute(C, head_view(C, qkv, head_dim, heads, pseq, kHidden, 3, 0), 0, 2, 1, 3)); - ggml_tensor * K = ggml_cont(C, ggml_permute(C, head_view(C, qkv, head_dim, heads, pseq, kHidden, 3, 1), 0, 2, 1, 3)); - ggml_tensor * V = ggml_cont(C, ggml_permute(C, head_view(C, qkv, head_dim, heads, pseq, kHidden, 3, 2), 1, 2, 0, 3)); + // fa_kv: on Hexagon K and V go to F16, so both matmuls run on the HMX + // rather than as F32 x F32 on the vector units. Scores accumulate in F32 + // and the softmax stays F32. Other backends keep F32 K/V. + ggml_tensor * K = fa_kv(C, ggml_cont(C, ggml_permute(C, head_view(C, qkv, head_dim, heads, pseq, kHidden, 3, 1), 0, 2, 1, 3))); + ggml_tensor * V = fa_kv(C, ggml_cont(C, ggml_permute(C, head_view(C, qkv, head_dim, heads, pseq, kHidden, 3, 2), 1, 2, 0, 3))); ggml_tensor * merged = attention(C, Q, K, V, mask, attn_scale, kHidden, pseq); ggml_tensor * residual = ggml_add(C, x, linear(C, blk_w[i][4], blk_w[i][5], merged)); diff --git a/src/models/openvla_oft.cpp b/src/models/openvla_oft.cpp index 4a7ecc77..9d460a5e 100644 --- a/src/models/openvla_oft.cpp +++ b/src/models/openvla_oft.cpp @@ -28,6 +28,7 @@ #include "layers/linear.h" #include "layers/norm.h" +#include #include #include #include @@ -68,6 +69,7 @@ struct OpenVlaOftModelArch : public ModelArchBase { }; struct MainIO { ggml_tensor *t_ids=nullptr,*t_state=nullptr,*t_proj=nullptr,*act0=nullptr,*t_pos=nullptr,*norm_actions=nullptr; + ggml_tensor *kv_mask=nullptr; bool consts=false; }; graph_cache main_graph; @@ -255,6 +257,9 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { seq=ggml_concat(C,seq,stop,1); ggml_tensor*t_pos=ggml_new_tensor_1d(C,GGML_TYPE_I32,SEQ); ggml_set_input(t_pos); ggml_set_output(t_pos); + const int64_t KP=gemm_width(SEQ); + ggml_tensor*kv_mask=nullptr; + if(KP>SEQ){ kv_mask=ggml_new_tensor_2d(C,GGML_TYPE_F32,KP,SEQ); ggml_set_input(kv_mask); ggml_set_output(kv_mask); } const float lsc=1.0f/std::sqrt((float)lm_head_dim); ggml_tensor*x=seq; const bool fused_lm=fused_ops && cublas_rows(SEQ) && SEQ<=4096; @@ -278,11 +283,21 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { ggml_tensor*qh=ggml_reshape_3d(C,qp,lm_head_dim,n_q,SEQ),*kh=ggml_reshape_3d(C,kp,lm_head_dim,n_kv,SEQ),*vh=ggml_reshape_3d(C,vp,lm_head_dim,n_kv,SEQ); ggml_tensor*qr=ggml_rope_ext(C,qh,t_pos,nullptr,(int)lm_head_dim,GGML_ROPE_TYPE_NEOX,0,lm_rope_base,1.0f,0.0f,1.0f,32.0f,1.0f); ggml_tensor*kr=ggml_rope_ext(C,kh,t_pos,nullptr,(int)lm_head_dim,GGML_ROPE_TYPE_NEOX,0,lm_rope_base,1.0f,0.0f,1.0f,32.0f,1.0f); + // fa_kv: on Hexagon K and V go to F16, and the keys are padded with zeros + // to gemm_width(SEQ), masked out, so both matmuls run on the HMX rather + // than as F32 x F32 on the vector units. Scores accumulate in F32 and the + // softmax stays F32. Other backends keep F32 K/V and SEQ keys. ggml_tensor*Q=ggml_cont(C,ggml_permute(C,qr,0,2,1,3)),*K=ggml_cont(C,ggml_permute(C,kr,0,2,1,3)),*V=ggml_cont(C,ggml_permute(C,vh,1,2,0,3)); + if(kv_mask){ + K=ggml_pad(C,K,0,(int)(KP-SEQ),0,0); + V=ggml_pad(C,V,(int)(KP-SEQ),0,0,0); + } + K=fa_kv(C,K); V=fa_kv(C,V); ggml_tensor*kq=ggml_mul_mat(C,K,Q); ggml_prec_set_acc(kq,GGML_PREC_F32); - // Unmasked on purpose: OpenVLA-OFT patches transformers to replace the - // causal mask across the whole sequence (modeling_llama.py:719-723). - ggml_tensor*aw=ggml_soft_max_ext(C,kq,nullptr,lsc,0.0f); + // No causal mask on purpose: OpenVLA-OFT patches transformers to replace + // it across the whole sequence (modeling_llama.py:719-723). kv_mask only + // hides the key padding. + ggml_tensor*aw=ggml_soft_max_ext(C,kq,kv_mask,lsc,0.0f); ggml_tensor*kqv=ggml_mul_mat(C,V,aw); ggml_tensor*att=ggml_reshape_2d(C,ggml_cont(C,ggml_permute(C,kqv,0,2,1,3)),HC,SEQ); x=ggml_add(C,x,ggml_mul_mat(C,l.Wo,att)); @@ -306,7 +321,7 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { ggml_tensor*norm_actions=linear(C,h_fc2w,h_fc2b,hh); ggml_set_output(norm_actions); gio.t_ids=t_ids; gio.t_state=t_state; gio.t_proj=t_proj; gio.act0=act0; - gio.t_pos=t_pos; gio.norm_actions=norm_actions; + gio.t_pos=t_pos; gio.norm_actions=norm_actions; gio.kv_mask=kv_mask; ggml_cgraph*gf=ggml_new_graph_custom(C,16384,false); ggml_build_forward_expand(gf,norm_actions); return gf; @@ -330,6 +345,13 @@ std::vector OpenVlaOftModelArch::predict(const Inputs& in) { for(int64_t i=0;ine[0]; + std::vector mk((size_t)(KP*SEQ),0.0f); + for(int64_t q=0;q z((size_t)HC*n_act,0.0f); ggml_backend_tensor_set(act0,z.data(),0,ggml_nbytes(act0)); gio.consts=true; diff --git a/src/models/vla_adapter.cpp b/src/models/vla_adapter.cpp index 25c65a13..b34a17ec 100644 --- a/src/models/vla_adapter.cpp +++ b/src/models/vla_adapter.cpp @@ -317,7 +317,10 @@ std::vector VlaAdapterModelArch::predict(const Inputs& in) { ggml_tensor*qh=ggml_reshape_3d(C,qp,lm_head_dim,n_q,LSEQ),*kh=ggml_reshape_3d(C,kp,lm_head_dim,n_kv,LSEQ),*vh=ggml_reshape_3d(C,vp,lm_head_dim,n_kv,LSEQ); ggml_tensor*qr=ggml_rope_ext(C,qh,t_pos,nullptr,(int)lm_head_dim,GGML_ROPE_TYPE_NEOX,0,lm_rope_base,1.0f,0.0f,1.0f,32.0f,1.0f); ggml_tensor*kr=ggml_rope_ext(C,kh,t_pos,nullptr,(int)lm_head_dim,GGML_ROPE_TYPE_NEOX,0,lm_rope_base,1.0f,0.0f,1.0f,32.0f,1.0f); - ggml_tensor*Q=ggml_cont(C,ggml_permute(C,qr,0,2,1,3)),*K=ggml_cont(C,ggml_permute(C,kr,0,2,1,3)),*V=ggml_cont(C,ggml_permute(C,vh,1,2,0,3)); + // fa_kv on V only: on Hexagon the PV matmul then runs on the HMX. K stays + // F32, since F16 scores put this LM 3.9e-3 from the CPU F32 reference on + // the IQ-10. Other backends keep F32 V. + ggml_tensor*Q=ggml_cont(C,ggml_permute(C,qr,0,2,1,3)),*K=ggml_cont(C,ggml_permute(C,kr,0,2,1,3)),*V=fa_kv(C,ggml_cont(C,ggml_permute(C,vh,1,2,0,3))); ggml_tensor*kq=ggml_mul_mat(C,K,Q); ggml_prec_set_acc(kq,GGML_PREC_F32); ggml_tensor*aw=ggml_soft_max_ext(C,kq,t_mask,lsc,0.0f); ggml_tensor*kqv=ggml_mul_mat(C,V,aw); diff --git a/src/models/vla_jepa.cpp b/src/models/vla_jepa.cpp index ff42d56d..f2c20073 100644 --- a/src/models/vla_jepa.cpp +++ b/src/models/vla_jepa.cpp @@ -148,6 +148,9 @@ bool load_config(const gguf_reader & g, const Options & opts, VlaJepaModelArch & m.dit.cfg.layers = m.dit_layers; m.dit.cfg.ln_eps = m.dit_ln_eps; m.dit.cfg.norm_out_eps = m.dit_norm_out_eps; + // Its cross-attention K/V overflow F16 on the IQ-10 (NaN actions); they + // stay F32 there, and only the self-attention takes the HMX. + m.dit.cfg.cross_kv_f16 = false; cfg = Config{}; cfg.n_img = m.vit.n_tokens(); diff --git a/src/modules/dit_head.cpp b/src/modules/dit_head.cpp index bb4250ad..1d37bb7e 100644 --- a/src/modules/dit_head.cpp +++ b/src/modules/dit_head.cpp @@ -152,8 +152,25 @@ void DitHead::declare(WeightLoader & L, const char * prefix, bool fuse_qkv, bool po2b = L.f32 ("%s.proj_out2.bias", outer); } +namespace { + +// K [hd, T, heads] and V^T [T, hd, heads] in F16, the head dim and the keys +// zero-padded to gemm_width(), as the HMX takes them (Hexagon). Elsewhere both +// come back unchanged. +void hmx_kv(ggml_context * C, ggml_tensor ** K, ggml_tensor ** V) { + const int64_t hd = (*K)->ne[0], T = (*K)->ne[1], HP = gemm_width(hd), NK = gemm_width(T); + if (HP > hd || NK > T) { + *K = ggml_pad(C, *K, (int) (HP-hd), (int) (NK-T), 0, 0); + *V = ggml_pad(C, *V, (int) (NK-T), (int) (HP-hd), 0, 0); + } + *K = fa_kv(C, *K); + *V = fa_kv(C, *V); +} + +} + void DitHead::kv(ggml_context * C, const DitLayerW & w, ggml_tensor * src, - ggml_tensor ** K_out, ggml_tensor ** V_out, ggml_tensor * xq_pre) const { + ggml_tensor ** K_out, ggml_tensor ** V_out, ggml_tensor * xq_pre, bool cross) const { const int64_t hd = cfg.head_dim; const int64_t heads = cfg.heads; const int64_t Tkv = src->ne[1]; @@ -187,6 +204,8 @@ void DitHead::kv(ggml_context * C, const DitLayerW & w, ggml_tensor * src, } *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)); + if (!cfg.fused_ops && (!cross || cfg.cross_kv_f16)) + hmx_kv(C, K_out, V_out); return; } if (w.fq_k) { @@ -198,6 +217,20 @@ void DitHead::kv(ggml_context * C, const DitLayerW & w, ggml_tensor * src, } *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); + if (!cfg.fused_ops && (!cross || cfg.cross_kv_f16)) + hmx_kv(C, K_out, V_out); +} + +ggml_tensor * DitHead::key_mask(ggml_context * C, int64_t NK, int64_t Nk, int64_t Tq) const { + char name[GGML_MAX_NAME]; + std::snprintf(name, sizeof(name), "dit.key_mask.%lld.%lld.%lld", (long long) NK, (long long) Nk, (long long) Tq); + if (ggml_tensor * m = ggml_get_tensor(C, name)) + return m; + // 0 for the Nk real keys, -1e30 for the padding: step(i - Nk + 0.5) * -1e30. + ggml_tensor * m = ggml_step(C, ggml_scale_bias(C, ggml_arange(C, 0.0f, (float) NK, 1.0f), 1.0f, 0.5f - (float) Nk)); + m = ggml_repeat_4d(C, ggml_scale(C, m, -1e30f), NK, Tq, 1, 1); + ggml_set_name(m, name); + return m; } ggml_tensor * DitHead::block(ggml_context * C, const DitLayerW & w, ggml_tensor * h, ggml_tensor * mod, @@ -278,6 +311,8 @@ ggml_tensor * DitHead::block(ggml_context * C, const DitLayerW & w, ggml_tensor 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)); + if (!cfg.fused_ops) + hmx_kv(C, &K, &V); } } else { // Self-attention with separate FoldQuant q/k/v: one activation blob of n @@ -294,7 +329,7 @@ ggml_tensor * DitHead::block(ggml_context * C, const DitLayerW & w, ggml_tensor V = V_pre; } else { - kv(C, w, enc ? enc : n, &K, &V, enc ? nullptr : xq); + kv(C, w, enc ? enc : n, &K, &V, enc ? nullptr : xq, enc != nullptr); } } @@ -310,8 +345,26 @@ ggml_tensor * DitHead::block(ggml_context * C, const DitLayerW & w, ggml_tensor ggml_tensor * g = fused_bias_gelu(C, mm_exact(C, w.Wff0, n3), w.bff0, bf16_operand(w.Wff2, Tk)); return fused_bias_res(C, mm_exact(C, w.Wff2, g), w.bff2, h1); } - ggml_tensor * att = exact ? fused_attn_exact(C, Q, K, V, scale, false) - : attention(C, Q, K, V, nullptr, scale, dim, Tk); + ggml_tensor * att; + // hmx_kv pads to gemm_width(), which is the identity off Hexagon, so these + // stay false there whatever K/V the caller passes. + const int64_t Nk = enc ? enc->ne[1] : Tk; + const bool hpad = K->ne[0] > hd && K->ne[0] == gemm_width(hd); + const bool kpad = K->ne[1] > Nk && K->ne[1] == gemm_width(Nk); + if (!exact && (hpad || kpad)) { + // hmx_kv padded K/V: the queries take the head-dim padding too, the + // padded keys are masked out and the output drops the padded columns. + if (hpad) + Q = ggml_pad(C, Q, (int) (K->ne[0]-hd), 0, 0, 0); + ggml_tensor * kq = ggml_mul_mat(C, K, Q); ggml_prec_set_acc(kq, GGML_PREC_F32); + ggml_tensor * mask = kpad ? key_mask(C, K->ne[1], Nk, Tk) : nullptr; + ggml_tensor * kqv = ggml_mul_mat(C, V, ggml_soft_max_ext(C, kq, mask, scale, 0.0f)); + kqv = ggml_view_3d(C, kqv, hd, Tk, heads, kqv->nb[1], kqv->nb[2], 0); + att = ggml_reshape_2d(C, ggml_cont(C, ggml_permute(C, kqv, 0, 2, 1, 3)), dim, Tk); + } else { + att = exact ? fused_attn_exact(C, Q, K, V, scale, false) + : attention(C, Q, K, V, nullptr, scale, dim, Tk); + } // 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 = cfg.fused_ops ? fused_ln(C, h1, nullptr, nullptr, cfg.ln_eps, false) : ggml_norm(C, h1, cfg.ln_eps); diff --git a/src/modules/dit_head.h b/src/modules/dit_head.h index 5eb573ed..cff7ff5c 100644 --- a/src/modules/dit_head.h +++ b/src/modules/dit_head.h @@ -50,6 +50,9 @@ struct DitCfg { float ln_eps = 1e-5f; float norm_out_eps = 1e-6f; bool fused_ops = false; // layers/fused_ops.h nodes (CUDA only) + // Cross-attention K/V in F16 where the HMX wants it (Hexagon). Off for a + // conditioning whose projections overflow F16. + bool cross_kv_f16 = true; }; struct DitHead { @@ -72,14 +75,18 @@ struct DitHead { const char * outer = nullptr, const FqModuleSpec * fq = nullptr); // xq_pre: an activation blob already computed from src (shared with q). + // cross: src is the conditioning rather than the block's own tokens. void kv(ggml_context * C, const DitLayerW & w, ggml_tensor * src, - ggml_tensor ** K_out, ggml_tensor ** V_out, ggml_tensor * xq_pre = nullptr) const; + ggml_tensor ** K_out, ggml_tensor ** V_out, ggml_tensor * xq_pre = nullptr, bool cross = true) const; ggml_tensor * block(ggml_context * C, const DitLayerW & w, ggml_tensor * h, ggml_tensor * mod, ggml_tensor * enc, ggml_tensor * K_pre = nullptr, ggml_tensor * V_pre = nullptr) const; ggml_tensor * time_emb(ggml_context * C, ggml_tensor * tproj) const; + // [NK, Tq] additive mask hiding keys Nk.. of NK, made once per graph (by name). + ggml_tensor * key_mask(ggml_context * C, int64_t NK, int64_t Nk, int64_t Tq) const; + // (shift, scale), opposite to the blocks. ggml_tensor * proj_out(ggml_context * C, ggml_tensor * h, ggml_tensor * mod) const; };