Static YaRN/linear RoPE long-context scaling + engine override (ABI v10)

CLAUDE.md

@@ -143,10 +143,21 @@ reference/.venv/bin/python reference/dump_ref.py
   still-unimplemented families (e.g. Unigram/WordPiece).
 - **Masks are additive fp16, never boolean** (avoids MLX bug #2894). See
   `DecoderModel::batch_mask`.
-- **RoPE is llama3-scaled** — `compute_rope_freqs` mirrors `mlx_lm`'s
-  `Llama3RoPE` exactly and feeds `fast::rope` via `freqs` (base disabled). It's
-  validated against `reference/fixtures/rope_freqs.npy`; don't "simplify" the
-  math.
+- **RoPE scaling is precomputed and validated, never silent** —
+  `compute_rope_setup` mirrors `mlx_lm` exactly (`Llama3RoPE`, `YarnRoPE`,
+  `nn.RoPE(scale=1/factor)` for linear) and feeds `fast::rope` via `freqs`
+  (base disabled); gated against `reference/fixtures/rope_freqs.npy` and
+  `reference/fixtures_qwen3_yarn/` — don't "simplify" the math. Two conventions
+  matter: yarn's attention **mscale multiplies the rope *input*** (both Q and
+  K, so logits scale by mscale² and the KV cache stores scaled K, exactly like
+  mlx-lm — folding it into the SDPA scale instead would desync the
+  q_rope0/k_rope0 fixtures), and **unknown `rope_type`s are rejected at load**
+  (`validate_rope_scaling`, run on the engine's caller thread — the worker
+  thread can't throw) rather than falling back to unscaled RoPE. The
+  engine-level override (`EngineConfig::rope_scaling`, ABI v10) fully replaces
+  the checkpoint's config and is re-applied inside `make_factory` on the
+  worker thread; keep the two applications identical or the factory can throw
+  where it must not.
 - **Decode-with-cache vs full-recompute logits differ by fp16 accumulation
   order** — compare argmax / exact tokens, not raw logits at tight tolerance.
 - **Quantized KV (kv_bits 8|4) mirrors mlx-lm's QuantizedKVCache** — triplet

apps/mlxforge.cpp

@@ -81,14 +81,16 @@ void print_help() {
       "      --prefill-chunk <N>  interleaved-prefill chunk size in tokens, 0 = monolithic\n"
       "                         (default 256: decode keeps streaming during prefills)\n"
       "      --skinny-mm <0|1>  multi-row GEMV decode kernels for small batches (default 1)\n"
+      "      --rope-scaling <J> RoPE-scaling JSON override (vLLM-style), e.g.\n"
+      "                         '{\"rope_type\":\"yarn\",\"factor\":4.0}' (default: checkpoint's)\n"
       "  -h, --help             show this help and exit\n"
       "\n"
       "The model may be given via -m or the config file's \"model\" key.\n"
       "Config precedence (low to high): defaults < config file < env vars < CLI flags.\n"
       "Env vars: MLXFORGE_HOST, MLXFORGE_PORT, MLXFORGE_MAX_CTX, MLXFORGE_MAX_WAITING, "
       "MLXFORGE_KV_BUDGET, MLXFORGE_KV_BITS, MLXFORGE_PREFIX_CACHE, MLXFORGE_KV_BLOCK, "
       "MLXFORGE_KV_POOL, MLXFORGE_KV_SPILL_DIR, MLXFORGE_KV_SPILL_BYTES, "
-      "MLXFORGE_PREFILL_CHUNK, MLXFORGE_SKINNY_MM.");
+      "MLXFORGE_PREFILL_CHUNK, MLXFORGE_SKINNY_MM, MLXFORGE_ROPE_SCALING.");
   std::fflush(stdout);
 }
 
@@ -148,6 +150,7 @@ int main(int argc, char** argv) {
     ec.kv_spill_bytes = sc.kv_spill_bytes;
     ec.prefill_chunk = sc.prefill_chunk;
     ec.skinny_mm = sc.skinny_mm;
+    ec.rope_scaling = sc.rope_scaling;
     engine = std::make_unique<mlxforge::Engine>(std::move(ec));
   } catch (const std::exception& e) {
     mlxforge::log::error("model error: {}", e.what());

apps/mlxforge_cli.cpp

@@ -7,11 +7,14 @@
 //     - Loads a model's weights from the supplied directory, prints key/shape/dtype for each tensor,
 //       asserts that all tensors are fp16, and reports the peak resident memory used.
 //   mlxforge-cli generate <model> <prompt> [max_tokens] [--logprobs [N]] [--kv-bits N]
+//                [--rope-scaling <json>]
 //     - Runs greedy single-stream generation: pre-fills the prompt (as raw text using the chat template
 //       or as a pre-tokenized .npy of ids), then streams the detokenized text to stdout until EOS or
 //       max_tokens. With --logprobs [N], each emitted token's log-prob (and its N most-likely
 //       alternatives) is printed to stderr after generation; stdout stays the generated text.
 //       --kv-bits 8|4 stores the KV cache quantized (the manual harness for the quantized path).
+//       --rope-scaling overrides the checkpoint's RoPE-scaling config with a JSON object
+//       (vLLM-style), e.g. '{"rope_type":"yarn","factor":4.0}' for long context.
 //   mlxforge-cli bench <model> [max_tokens] [runs]
 //     - Repeatable throughput benchmark over a fixed prompt: one discarded warmup run, then `runs`
 //       timed runs (defaults: max_tokens=128, runs=3) reporting time-to-first-token and decode tok/s.
@@ -77,18 +80,24 @@ struct LoadedModel {
 };
 
 // Resolve a model spec and load it for single-stream inference, dispatching on
-// whether it resolves to a GGUF file or a safetensors directory.
-LoadedModel load_for_inference(const std::string& spec) {
+// whether it resolves to a GGUF file or a safetensors directory. `rope_scaling`
+// optionally overrides the checkpoint's RoPE-scaling config (the CLI bypasses
+// Engine, so the override is applied here; model construction validates it).
+LoadedModel load_for_inference(const std::string& spec, const std::string& rope_scaling = "") {
   const std::string resolved = mlxforge::resolve_model_dir(spec);
   LoadedModel lm;
   if (mlxforge::is_gguf_path(resolved)) {
+    if (!rope_scaling.empty())
+      throw std::runtime_error("rope_scaling override is not supported for GGUF models");
     mlxforge::GgufModel g = mlxforge::load_gguf_model(resolved);
     lm.cfg = g.config;
     lm.tok = mlxforge::Tokenizer::from_gguf(g.tokens, g.merges, g.token_types, g.pre, g.bos_id,
                                             mlxforge::chat_format_from_model_type(lm.cfg.model_type));
     lm.model = mlxforge::create_model(std::move(g.config), std::move(g.weights));
   } else {
     lm.cfg = mlxforge::ModelConfig::from_file(resolved + "/config.json");
+    if (!rope_scaling.empty()) mlxforge::apply_rope_scaling_override(lm.cfg, rope_scaling);
+    mlxforge::validate_rope_scaling(lm.cfg);
     lm.model = mlxforge::create_model(lm.cfg, mlxforge::load_weights(resolved, lm.cfg));
     lm.tok = mlxforge::Tokenizer::from_file(resolved + "/tokenizer.json", lm.cfg.bos_token_id,
                                             mlxforge::chat_format_from_model_type(lm.cfg.model_type));
@@ -188,9 +197,10 @@ std::string show_token(const std::string& s) {
 // `top_logprobs` mirrors the engine knob: -1 = off; 0 = each token's own log-prob;
 // N > 0 = also its N most-likely alternatives (printed to stderr after generation).
 int run_generate(const std::string& spec, const std::string& prompt_arg, int max_tokens,
-                 int top_logprobs = -1, int kv_bits = 0) {
+                 int top_logprobs = -1, int kv_bits = 0,
+                 const std::string& rope_scaling = "") {
   // Resolve and load the model (GGUF file or safetensors dir; downloads if needed)
-  LoadedModel lm = load_for_inference(spec);
+  LoadedModel lm = load_for_inference(spec, rope_scaling);
   mlxforge::DecoderModel& model = *lm.model;
   mlxforge::Tokenizer& tok = lm.tok;
   const mlxforge::ModelConfig& cfg = lm.cfg;
@@ -466,15 +476,17 @@ int main(int argc, char** argv) {
     if (argc < 4) {
       std::fprintf(stderr,
                    "usage: mlxforge-cli generate <model_dir> <prompt_ids.npy> [max_tokens] "
-                   "[--logprobs [N]] [--kv-bits N]\n");
+                   "[--logprobs [N]] [--kv-bits N] [--rope-scaling <json>]\n");
       return 2;
     }
     // Positional [max_tokens] (default 64), an optional --logprobs [N] flag (N
-    // alternatives, default 0 = the chosen token's own log-prob only), and an
-    // optional --kv-bits N (0 = fp16 cache, 8 or 4 = quantized).
+    // alternatives, default 0 = the chosen token's own log-prob only), an
+    // optional --kv-bits N (0 = fp16 cache, 8 or 4 = quantized), and an optional
+    // --rope-scaling JSON override (vLLM-style, e.g. yarn for long context).
     int max_tokens = 64;
     int top_logprobs = -1;
     int kv_bits = 0;
+    std::string rope_scaling;
     for (int i = 4; i < argc; ++i) {
       const std::string a = argv[i];
       if (a == "--logprobs") {
@@ -492,11 +504,17 @@ int main(int argc, char** argv) {
           std::fprintf(stderr, "error: --kv-bits must be 0, 4, or 8\n");
           return 2;
         }
+      } else if (a == "--rope-scaling") {
+        if (i + 1 >= argc) {
+          std::fprintf(stderr, "error: --rope-scaling needs a JSON value\n");
+          return 2;
+        }
+        rope_scaling = argv[++i];
       } else {
         max_tokens = std::stoi(a);
       }
     }
-    return run_generate(argv[2], argv[3], max_tokens, top_logprobs, kv_bits);
+    return run_generate(argv[2], argv[3], max_tokens, top_logprobs, kv_bits, rope_scaling);
   }
   if (cmd == "image") {
     // Vision-language generation: describe / answer about an image.

doc/supported-models.md

@@ -33,6 +33,18 @@ the **ChatML** chat template (with an `enable_thinking` toggle for Qwen3's
 reasoning mode), and single-digit number pre-tokenization in the byte-level BPE.
 Qwen3 has **no BOS token**.
 
+**RoPE scaling.** Safetensors checkpoints may ship a `rope_scaling` config:
+`llama3` (Llama-3.2's rescaling), **`yarn`** and **`linear`** (long context — e.g.
+Qwen3's documented yarn recipe takes its native 32k window to 131k) are
+implemented and golden-gated against mlx-lm (`reference/fixtures_qwen3_yarn/`);
+`default` means unscaled. Yarn/linear can also be *forced onto a stock
+checkpoint* via the engine-level override (`mlxforge_engine_opts2.rope_scaling`,
+the server/CLI `--rope-scaling` flag), e.g.
+`'{"rope_type":"yarn","factor":4.0}'`. Anything else — unknown types
+(`dynamic`, `longrope`), yarn/linear on hybrid (Qwen3.5) or vision (Qwen3-VL)
+models, or on GGUF checkpoints — **fails at load** with a clear error; there is
+never a silent fall-back to unscaled RoPE.
+
 **Qwen3 MoE** models (e.g. 30B-A3B, 235B-A22B) run end-to-end too. They share the
 dense Qwen3 attention (QK-Norm) and ChatML tokenizer; the only delta is the
 feed-forward block. On the MoE layers (selected by `config.json`'s `num_experts`,

reference/dump_ref.py

@@ -22,6 +22,7 @@
 import os
 
 import mlx.core as mx
+import mlx.nn as nn
 import numpy as np
 from mlx_lm import load
 from mlx_lm.models.base import create_attention_mask, create_ssm_mask
@@ -49,6 +50,26 @@
         "chat_messages": [{"role": "user", "content": "What is the capital of France?"}],
         "thinking": True,
     },
+    # YaRN long-context gate: the SAME Qwen3 weights with Qwen's documented yarn
+    # recipe injected via model_config (no checkpoint on the Hub ships yarn in
+    # config.json, so it is injected on both sides — the C++ test reads the exact
+    # object back from this manifest and applies it through the engine's override
+    # path). Gates the yarn freqs schedule, the mscale-on-input convention, and
+    # the scaled forward end to end against mlx_lm's YarnRoPE.
+    "qwen3_yarn": {
+        "repo": "mlx-community/Qwen3-0.6B-bf16",
+        "fixtures": "fixtures_qwen3_yarn",
+        "compute_dtype": mx.float16,
+        "chat_messages": [{"role": "user", "content": "What is the capital of France?"}],
+        "thinking": True,
+        "model_config": {
+            "rope_scaling": {
+                "rope_type": "yarn",
+                "factor": 4.0,
+                "original_max_position_embeddings": 32768,
+            }
+        },
+    },
     # Qwen3 MoE shares the dense Qwen3 attention (QK-Norm) and ChatML tokenizer;
     # only the feed-forward block differs (sparse expert routing). The smallest
     # MoE checkpoint is 30B-A3B; the 4bit repo keeps the download manageable.
@@ -536,7 +557,10 @@ def main():
 
     os.makedirs(FIXTURES_DIR, exist_ok=True)
     print(f"loading {MODEL_REPO} ...")
-    model, tok = load(MODEL_REPO)
+    # model_config (e.g. an injected yarn rope_scaling) updates config.json before
+    # the model is built; the manifest records it so the C++ test injects the
+    # exact same object through the engine's override path.
+    model, tok = load(MODEL_REPO, model_config=spec.get("model_config"))
     if COMPUTE_DTYPE is not None:
         model.set_dtype(COMPUTE_DTYPE)
     mx.eval(model.parameters())
@@ -549,6 +573,8 @@ def main():
         "greedy_max_new": GREEDY_MAX_NEW,
         "arrays": {},
     }
+    if spec.get("model_config"):
+        manifest["model_config"] = spec["model_config"]
 
     def save(name, arr):
         """Eval an MLX array (or accept a numpy array) and write it as .npy."""
@@ -627,15 +653,30 @@ def save(name, arr):
         k = attn.k_norm(k)
     q = q.transpose(0, 2, 1, 3)  # (1, n_heads, T, head_dim)
     k = k.transpose(0, 2, 1, 3)  # (1, n_kv_heads, T, head_dim)
-    # The precomputed `_freqs` is specific to the llama3-rescaled RoPE; plain-RoPE
-    # models (Qwen3) lack it. The front-half intermediates are dumped for both.
+    # The precomputed `_freqs` is specific to rescaled RoPE (llama3, yarn);
+    # plain-RoPE models (Qwen3) lack it. The front-half intermediates are dumped
+    # for both. YarnRoPE additionally carries the attention mscale it folds into
+    # the rope input — dumped so the C++ side gates the same convention.
     if hasattr(attn.rope, "_freqs"):
-        save("rope_freqs", attn.rope._freqs)  # (head_dim/2,) llama3-rescaled freqs
+        save("rope_freqs", attn.rope._freqs)  # (head_dim/2,) rescaled freqs
+    if hasattr(attn.rope, "mscale"):
+        save("rope_mscale", np.float32(attn.rope.mscale))
     save("q_pre0", q)  # pre-RoPE queries (post q_norm for Qwen3)
     save("q_rope0", attn.rope(q))
     save("k_rope0", attn.rope(k))
     save("v0", v)
 
+    # Linear-scaling oracle, piggybacked on the yarn dump (no separate fixture
+    # dir): mlx_lm's rope_type "linear" is nn.RoPE(scale=1/factor); apply it to
+    # the same committed q_pre0/k_pre0 tensors so the C++ linear freqs path has a
+    # true fixture-to-fixture gate.
+    if spec.get("model_config", {}).get("rope_scaling", {}).get("rope_type") == "yarn":
+        factor = spec["model_config"]["rope_scaling"]["factor"]
+        lin = nn.RoPE(q.shape[-1], traditional=False, base=model.args.rope_theta,
+                      scale=1.0 / factor)
+        save("q_rope0_linear", lin(q))
+        save("k_rope0_linear", lin(k))
+
     # Block-0 output: exactly what LlamaModel.__call__ computes for layer 0
     # (gates the single decoder block).
     mask = create_attention_mask(embeddings, cache=None)