Pull requests / #2819
#2819 metal: --rope-scaling, --rope-scale, --rope-freq-base, --rope-freq-scale and --yarn-* reach the decode window's RoPE table
open · @oinani0721 · 0 comentários · No GitHub
BenchmarksSetup & installAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindowsLinux
Descrição
## Title
Issue: Resolves #2767. No other issue or pull request covers RoPE scaling on the Mac engine (search in the first comment).
Every file:line is at 296b59d (v0.1.44.1, the last release) unless marked "branch". Short names are the files under src/, include/ or metal/, not their copies under sycl/; CMakeLists.txt is the root one. On the branch, window.cpp's lines 18-467 move 1 down and its lines after 475 move 4 up, setup_host.cpp's lines after 125 move 6 down, and metal/CMakeLists.txt's lines 106-123 move 9 down and its lines after 123 move 16 down.
## Summary
On a Mac, `--rope-scaling yarn --rope-scale 4` is parsed, resolved and logged as active (`strata generate: rope scaling yarn, factor 4 ...`, generate.cpp:3002-3008), but, by the code, no rotation changes:
- `rope_table_set` does nothing on Metal (metal/src/engine/setup_host.cpp:151, "the window keeps its own table");
- every RoPE rotation on Metal is in `qsa_prep`, which reads only the decode window's own cos/sin table (qsa.metal:74);
- `Window::init` builds that table without scaling, from the model file's `rope.freq_base` (window.cpp:468-475).
This PR makes the window build the table with `build_rope_table(64, rope_scaling(), ...)`, the builder the PC's session table uses (layer.cpp:837-840). setup_host.cpp:116-150 keeps rope.cu:48-86: the same 417 tokens once comments, whitespace and braces are ignored. By the code, linear and YaRN scaling, the magnitude factor (YaRN's mscale, or `--yarn-attn-factor`), `--rope-freq-base` and `--rope-freq-scale` then reach every rotation on a Mac without a kernel change (the reading is in the first comment). On an M5 Max (On-model results), yarn x4 changes the decode windows' logits on this branch and changes nothing on main (#2767); for each prompt path (the resident prompt chain, the streaming prompt GEMM) the test compares how close the path stays under yarn x4 to the same prompts read through decode windows, against a threshold derived in the same sitting, and for both paths that threshold did not separate the measured numbers, so they stay a code reading, as the MTP draft layer does.
This is one engine part needed to run past the trained 262,144 tokens on a Mac, not all of them: Mac setup does not write the flags, and nothing checks the memory a long context takes on a resident pack.
## What changed
- `metal/src/engine/window.cpp`: the table loop in `Window::init` becomes one `window_rope_table(strata::kernels::rope_scaling(), MAXC + 16, cs_.as<float>())` call. The table keeps its size (MAXC + 16 rows) and layout (float2 [pos][32]).
- `metal/src/engine/rope_window.h` (new): `window_rope_table` builds the table with `build_rope_table` into two temporary vectors of (MAXC + 16) x 32 floats each (branch rope_window.h:21), freed when it returns, and interleaves them into the window's rows; nothing it allocates stays. The window and the test run this same code.
- `metal/src/engine/setup_host.cpp`: a declaration of the scaled `build_rope_table` with `__attribute__((flatten))` (branch :126-131), so `RopeScaling::mscale()` is inlined into it under the flag below; the function's body does not change. The file's opening comment now says that `rope_scaling()` and `build_rope_table` fill the window's table.
- `metal/CMakeLists.txt`: `set_source_files_properties(${CMAKE_CURRENT_SOURCE_DIR}/src/engine/setup_host.cpp PROPERTIES COMPILE_OPTIONS "-ffp-contract=off")`, the form CMakeLists.txt:450-453 uses for two HIP files, with a comment that gives the by-hand check for the flatten (below); and the `rope_table_parity` target with `add_test`. The test has its own folder because `tests/*.cpp` are globbed into mtl_rt-only programs (metal/CMakeLists.txt:41-47). Setup builds only `strata` (setup_mac.py:280).
- `metal/tests/engine/rope_table_parity.cpp` (new; CPU only, no GPU or model; 5.7 s and 0.54 GB in ctest on the M5 Max). At the allocation of the 1M probe (bench/results/2026-10-09-rope-yarn-long-context/probe.py, whose default capacity is `--max-context 1048832`, probe.py:32), 1,048,976 rows, it checks:
1. unscaled: the new table against the replaced loop, byte for byte at every row; and with `--rope-freq-base 5e6`, against that loop given the same base (a copy of the test whose table ignores the base fails here);
2. yarn x2, x3 and x4, linear x4, yarn x4 with `--yarn-attn-factor 0.9`, linear x4 with `--yarn-attn-factor 1.25` and yarn with `--rope-freq-scale 0.3`: bit for bit against a float64 transcription of ggml's `rope_yarn` (llama.cpp 3cf03257, ggml-cpu/ops.cpp:5958-5973) with every operation rounded on its own. As in src/kernels/rope_parity.cpp:233-236, the correction dims and the ramp come from the shared helpers; an oracle written from ggml's own formulas checks those too (first comment, with its code);
3. row 0 is (mscale, 0); under YaRN the correction window is pairs [14, 22], and at the last row the first pair extrapolates and the last pair interpolates.
What it leaves alone:
- **The Metal kernels.** No `.metal` file is touched, the spare indexer key included (Extra Notes).
- **Mac setup.** On a Mac, setup.py hands every argument to `metal/setup_mac.py` before its own parser runs (setup.py:5706-5707, 5446), so the `--rope-scaling` and `--rope-scale` it defines (setup.py:5715-5719) never reach that parser; setup_mac.py still drops them with "not used on a Mac" (setup_mac.py:508-510) and writes no rope flags. Setup support is a later PR, which has to count a long context's memory first.
- **Every file outside `metal/`.** The root CMakeLists.txt and sycl/CMakeLists.txt do not reference `metal/`, so no CUDA, HIP or SYCL build compiles a changed file. They were not built.
- **docs/MAC.md.** "Context: measured up to 32K" (MAC.md:176) changes only with measured results.
### How the default path was checked
All on an M5 Max (128 GB), AppleClang 21.0.0, against the Metal build of 296b59d; the commands and their output lines are in the first comment.
1. **The table.** Test check 1: byte-identical at all 1,048,976 rows; a row depends only on its position, so every smaller context's table is a prefix. On the default path `rope_scaling()` is `RopeScaling{}` with `freq_base` from the model file's `qwen4exp.rope.freq_base` (generate.cpp:2940, 2969, 2978-2990), the key the window read before.
2. **The model files.** The GGUF headers (metadata only) of the six models Mac setup offers (setup_mac.py:45-55) and of the abliterated IQ3_S the on-model tests use all carry `qwen4exp.rope.freq_base = 10000000.0` and no `rope.scaling.*` key.
3. **The object code.** Of the 70 objects both builds have, `otool -X -s __TEXT __text <object> | shasum` differs only for `setup_host.cpp.o` (97573d5fa328b5c1 -> fd4ef0c308696cb7) and `window.cpp.o` (afaf2e316677789d -> fc9d170bb55b5ea8). In setup_host.cpp.o 36 of 37 functions are identical; only `build_rope_table(int, const RopeScaling&, int, float*, float*)` changes. In window.cpp.o, with the checkout path normalized, only `Window::init` changes; `window_rope_table` is new and the `std::map` lookup the old loop made is gone.
4. **Build and tests.** ctest passes 2/2 (`runtime_hello`, `rope_table_parity`). `-Wall -Wextra` gives no warning on setup_host.cpp, the test or a file that includes only rope_window.h. `tools.test_setup_mac` and `tools.test_setup_rope` pass (38 tests); no setup file is changed.
### Why `-ffp-contract=off`, and why `flatten`
The project holds this table bit for bit to a float64 reference (rope_scaling.hpp:24-26, 122-123). Without the flag Apple clang 21 fuses YaRN's blend and mscale (4 fused multiply-adds), and at 1M 2 of the 67,134,464 floats of the yarn x2 table differ from the reference (sin at position 589,449 pair 19 and at 910,519 pair 17); at `--max-context` 262,144 and 524,544 none do, which is why ctest runs the 1M allocation. With the flag all seven scaled configurations match.
`generate.cpp.o` keeps an out-of-line `RopeScaling::mscale()` with a fused multiply-add. Compiled with `-fno-inline` as a stand-in for a compiler that stops inlining it, the builder without `flatten` calls `mscale()`, and a strata link binds that call to generate.cpp.o's copy; the flattened builder makes no such call. Today's build inlines it either way, and the flattened build gives the same eight tables (hashes in the first comment). ctest cannot see a lost `flatten`, since the test does not link generate.cpp.o; the CMake comment gives an `nm` check, for the default Release build. With only mscale fused, the 1M yarn table would change at 1 of the 252 factors from 1.25 to 64 in steps of 0.25 (36.25, in 1 float).
Inferred from the build files, not checked on a PC: rope.cu gets no `-march`, `-mfma` or `/arch:` (CMakeLists.txt:354, 1219-1225, 1326-1352), so a PC whose host compiler targets baseline x86-64 rounds it the same way. A host compiler that can fuse by default (an ARM64 Linux host, or `-march=native` in the flags) may not.
## On-model results
Apple M5 Max (128 GB, macOS 27.0.1), 2026-10-16; Qwen3.8-Flash-Next-GSQ-RCO-abliterated-IQ3_S (revision eec4e10a, hashes in the first comment), greedy, thinking off; main is 296b59d's Metal build. B0-B2 and B6 ran resident at `--max-context 65536`, except B2's two streaming runs; B3-B5 streamed with `--kv int8`. Desktop programs stayed open (30-42 GB in use beside the engine), so every streaming start ran with its cap lowered by `STRATA_METAL_LEAVE_SHARE` (recorded per run) and with the model's pages dropped from the file cache first; the first comment says why and has every table.
- **B1, the default path (no flag): main against this branch.** One-token decode windows, 86 requests: tokens, DONE fields and draft counts identical, 86 of 86 first-window logit files and 191 of 191 `STRATA_METAL_LOGPOS` rows byte-identical. Setup's default config, 6 requests (decided: the control, this branch run twice, was identical): tokens and DONE fields identical, 6 of 6 logit files byte-identical.
- **B2 on main (#2767):** yarn x4 (`--rope-scaling yarn --rope-scale 4 --yarn-orig-ctx 262144`) against no flag gives the same results in every comparison of B1, while the log says `rope scaling yarn, factor 4`.
- **B2 on this branch:** yarn x4 changes all 6 first-window rows of the default config and the logits at all 65 decode-window positions (mean KL 0.0152, top-1 agreement 0.923), and after all 20 prompts in the resident (mean KL 0.0276) and the streaming runs (0.0343). The prompt-path gate is undecided for both paths: the resident prompt chain n 0.0105, e 0.0318, d 0.0084, T 0.0183; the streaming prompt GEMM n 0.0070, e 0.0318, d 0.0101, T 0.0149. d is under T on the means, but T did not separate every prompt's n and e.
- **B6, needles at 8K, 16K and 32K:** 9 of 9 on main, on this branch without a flag (twice: the control was identical) and with yarn x4; main against this branch identical.
- **B3 and B3M, memory (streaming):** the engine holds 18.00, 25.62 and 41.08 GB before its arena at 262,144, 524,544 and 1,048,832 cells, with prompt buffers of 5.19, 6.26 and 8.41 GB; the least a 1M start can hold, computed with the smallest arena the harness accepts, is 58.5 GB, 24.0 GB under the engine's cap of 82.5 GB. `--kv int8` against fp16 at 32K: identical tokens and byte-identical logits, as #2707 reads the code (`qsa_prep` stores K and V as int8 whatever `--kv` says).
- **B4, 512K** (probe.py's three-code prompt, 524,288 tokens, the PC run's input by SHA-256): without scaling 3 of 3, prefill 752 tok/s (697 s), decode 22.1 tok/s; yarn x2 3 of 3, 657 tok/s (799 s), 21.6 tok/s. The cap was lowered to 55.9 GB (a 29.7 GB arena), so prompts read experts again more often than with the engine's own arena.
- **B5, 1M** (probe.py's three-code prompt, 1,048,576 tokens, the PC run's input by SHA-256; `STRATA_WATCHDOG_S=300`, the harness's setting when the arena is smaller than B4's): without scaling 1 of 3 (`CEDAR-482|MARBLE-956|QUARTZ-956`), prefill 589 tok/s (1,781 s), decode 20.0 tok/s; yarn x4 3 of 3, 580 tok/s (1,807 s), No site
Links install, modelos, releases.