diff --git a/CLAUDE.md b/CLAUDE.md index 2191a4e..3b178a9 100644 --- a/CLAUDE.md +++ b/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 diff --git a/apps/mlxforge.cpp b/apps/mlxforge.cpp index 36b8afc..3f330b4 100644 --- a/apps/mlxforge.cpp +++ b/apps/mlxforge.cpp @@ -81,6 +81,8 @@ void print_help() { " --prefill-chunk 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 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" @@ -88,7 +90,7 @@ void print_help() { "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(std::move(ec)); } catch (const std::exception& e) { mlxforge::log::error("model error: {}", e.what()); diff --git a/apps/mlxforge_cli.cpp b/apps/mlxforge_cli.cpp index 2fbdfca..57e17d5 100644 --- a/apps/mlxforge_cli.cpp +++ b/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 [max_tokens] [--logprobs [N]] [--kv-bits N] +// [--rope-scaling ] // - 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 [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,11 +80,15 @@ 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, @@ -89,6 +96,8 @@ LoadedModel load_for_inference(const std::string& spec) { 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 [max_tokens] " - "[--logprobs [N]] [--kv-bits N]\n"); + "[--logprobs [N]] [--kv-bits N] [--rope-scaling ]\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. diff --git a/doc/supported-models.md b/doc/supported-models.md index d36b666..812f74b 100644 --- a/doc/supported-models.md +++ b/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`, diff --git a/reference/dump_ref.py b/reference/dump_ref.py index 543f9b6..11d8080 100644 --- a/reference/dump_ref.py +++ b/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) diff --git a/reference/fixtures_qwen3_yarn/argmax.npy b/reference/fixtures_qwen3_yarn/argmax.npy new file mode 100644 index 0000000..2c78aa4 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/argmax.npy differ diff --git a/reference/fixtures_qwen3_yarn/attn_norm0.npy b/reference/fixtures_qwen3_yarn/attn_norm0.npy new file mode 100644 index 0000000..6feb9fe Binary files /dev/null and b/reference/fixtures_qwen3_yarn/attn_norm0.npy differ diff --git a/reference/fixtures_qwen3_yarn/block0.npy b/reference/fixtures_qwen3_yarn/block0.npy new file mode 100644 index 0000000..a59a7ab Binary files /dev/null and b/reference/fixtures_qwen3_yarn/block0.npy differ diff --git a/reference/fixtures_qwen3_yarn/chat_ids.npy b/reference/fixtures_qwen3_yarn/chat_ids.npy new file mode 100644 index 0000000..1a84f69 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/chat_ids.npy differ diff --git a/reference/fixtures_qwen3_yarn/chat_ids_nothink.npy b/reference/fixtures_qwen3_yarn/chat_ids_nothink.npy new file mode 100644 index 0000000..44b6a3e Binary files /dev/null and b/reference/fixtures_qwen3_yarn/chat_ids_nothink.npy differ diff --git a/reference/fixtures_qwen3_yarn/embeddings.npy b/reference/fixtures_qwen3_yarn/embeddings.npy new file mode 100644 index 0000000..1ce63b8 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/embeddings.npy differ diff --git a/reference/fixtures_qwen3_yarn/greedy_gaps_kvq4.npy b/reference/fixtures_qwen3_yarn/greedy_gaps_kvq4.npy new file mode 100644 index 0000000..d52d791 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/greedy_gaps_kvq4.npy differ diff --git a/reference/fixtures_qwen3_yarn/greedy_gaps_kvq8.npy b/reference/fixtures_qwen3_yarn/greedy_gaps_kvq8.npy new file mode 100644 index 0000000..29d5b5a Binary files /dev/null and b/reference/fixtures_qwen3_yarn/greedy_gaps_kvq8.npy differ diff --git a/reference/fixtures_qwen3_yarn/greedy_tokens.npy b/reference/fixtures_qwen3_yarn/greedy_tokens.npy new file mode 100644 index 0000000..b60e9cf Binary files /dev/null and b/reference/fixtures_qwen3_yarn/greedy_tokens.npy differ diff --git a/reference/fixtures_qwen3_yarn/greedy_tokens_kvq4.npy b/reference/fixtures_qwen3_yarn/greedy_tokens_kvq4.npy new file mode 100644 index 0000000..86f51a6 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/greedy_tokens_kvq4.npy differ diff --git a/reference/fixtures_qwen3_yarn/greedy_tokens_kvq8.npy b/reference/fixtures_qwen3_yarn/greedy_tokens_kvq8.npy new file mode 100644 index 0000000..2f62557 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/greedy_tokens_kvq8.npy differ diff --git a/reference/fixtures_qwen3_yarn/k_rope0.npy b/reference/fixtures_qwen3_yarn/k_rope0.npy new file mode 100644 index 0000000..2ac3071 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/k_rope0.npy differ diff --git a/reference/fixtures_qwen3_yarn/k_rope0_linear.npy b/reference/fixtures_qwen3_yarn/k_rope0_linear.npy new file mode 100644 index 0000000..6ca7408 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/k_rope0_linear.npy differ diff --git a/reference/fixtures_qwen3_yarn/logits_last.npy b/reference/fixtures_qwen3_yarn/logits_last.npy new file mode 100644 index 0000000..178360c Binary files /dev/null and b/reference/fixtures_qwen3_yarn/logits_last.npy differ diff --git a/reference/fixtures_qwen3_yarn/manifest.json b/reference/fixtures_qwen3_yarn/manifest.json new file mode 100644 index 0000000..825697a --- /dev/null +++ b/reference/fixtures_qwen3_yarn/manifest.json @@ -0,0 +1,189 @@ +{ + "model_repo": "mlx-community/Qwen3-0.6B-bf16", + "compute_dtype": "float16", + "prompts": [ + "The capital of France is", + "Hello, world!", + "Once upon a time, in a land far away," + ], + "chat_messages": [ + { + "role": "user", + "content": "What is the capital of France?" + } + ], + "greedy_max_new": 20, + "arrays": { + "prompt_0_ids": { + "shape": [ + 5 + ], + "dtype": "int32" + }, + "prompt_1_ids": { + "shape": [ + 4 + ], + "dtype": "int32" + }, + "prompt_2_ids": { + "shape": [ + 11 + ], + "dtype": "int32" + }, + "chat_ids": { + "shape": [ + 15 + ], + "dtype": "int32" + }, + "chat_ids_nothink": { + "shape": [ + 19 + ], + "dtype": "int32" + }, + "embeddings": { + "shape": [ + 1, + 5, + 1024 + ], + "dtype": "float16" + }, + "attn_norm0": { + "shape": [ + 1, + 5, + 1024 + ], + "dtype": "float16" + }, + "rope_freqs": { + "shape": [ + 64 + ], + "dtype": "float32" + }, + "rope_mscale": { + "shape": [], + "dtype": "float32" + }, + "q_pre0": { + "shape": [ + 1, + 16, + 5, + 128 + ], + "dtype": "float16" + }, + "q_rope0": { + "shape": [ + 1, + 16, + 5, + 128 + ], + "dtype": "float16" + }, + "k_rope0": { + "shape": [ + 1, + 8, + 5, + 128 + ], + "dtype": "float16" + }, + "v0": { + "shape": [ + 1, + 8, + 5, + 128 + ], + "dtype": "float16" + }, + "q_rope0_linear": { + "shape": [ + 1, + 16, + 5, + 128 + ], + "dtype": "float16" + }, + "k_rope0_linear": { + "shape": [ + 1, + 8, + 5, + 128 + ], + "dtype": "float16" + }, + "block0": { + "shape": [ + 1, + 5, + 1024 + ], + "dtype": "float16" + }, + "logits_last": { + "shape": [ + 1, + 151936 + ], + "dtype": "float16" + }, + "argmax": { + "shape": [ + 1 + ], + "dtype": "int32" + }, + "greedy_tokens": { + "shape": [ + 20 + ], + "dtype": "int32" + }, + "greedy_tokens_kvq8": { + "shape": [ + 20 + ], + "dtype": "int32" + }, + "greedy_gaps_kvq8": { + "shape": [ + 20 + ], + "dtype": "float32" + }, + "greedy_tokens_kvq4": { + "shape": [ + 20 + ], + "dtype": "int32" + }, + "greedy_gaps_kvq4": { + "shape": [ + 20 + ], + "dtype": "float32" + } + }, + "model_config": { + "rope_scaling": { + "rope_type": "yarn", + "factor": 4.0, + "original_max_position_embeddings": 32768 + } + }, + "eos_token_ids": [ + 151645 + ] +} \ No newline at end of file diff --git a/reference/fixtures_qwen3_yarn/prompt_0_ids.npy b/reference/fixtures_qwen3_yarn/prompt_0_ids.npy new file mode 100644 index 0000000..fa2ad06 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/prompt_0_ids.npy differ diff --git a/reference/fixtures_qwen3_yarn/prompt_1_ids.npy b/reference/fixtures_qwen3_yarn/prompt_1_ids.npy new file mode 100644 index 0000000..b6d2e9a Binary files /dev/null and b/reference/fixtures_qwen3_yarn/prompt_1_ids.npy differ diff --git a/reference/fixtures_qwen3_yarn/prompt_2_ids.npy b/reference/fixtures_qwen3_yarn/prompt_2_ids.npy new file mode 100644 index 0000000..4ce7600 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/prompt_2_ids.npy differ diff --git a/reference/fixtures_qwen3_yarn/q_pre0.npy b/reference/fixtures_qwen3_yarn/q_pre0.npy new file mode 100644 index 0000000..418ec44 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/q_pre0.npy differ diff --git a/reference/fixtures_qwen3_yarn/q_rope0.npy b/reference/fixtures_qwen3_yarn/q_rope0.npy new file mode 100644 index 0000000..b30a579 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/q_rope0.npy differ diff --git a/reference/fixtures_qwen3_yarn/q_rope0_linear.npy b/reference/fixtures_qwen3_yarn/q_rope0_linear.npy new file mode 100644 index 0000000..de9b397 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/q_rope0_linear.npy differ diff --git a/reference/fixtures_qwen3_yarn/rope_freqs.npy b/reference/fixtures_qwen3_yarn/rope_freqs.npy new file mode 100644 index 0000000..4ed35ff Binary files /dev/null and b/reference/fixtures_qwen3_yarn/rope_freqs.npy differ diff --git a/reference/fixtures_qwen3_yarn/rope_mscale.npy b/reference/fixtures_qwen3_yarn/rope_mscale.npy new file mode 100644 index 0000000..7d48bf6 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/rope_mscale.npy differ diff --git a/reference/fixtures_qwen3_yarn/tokenizer_corpus.json b/reference/fixtures_qwen3_yarn/tokenizer_corpus.json new file mode 100644 index 0000000..d949f7e --- /dev/null +++ b/reference/fixtures_qwen3_yarn/tokenizer_corpus.json @@ -0,0 +1,392 @@ +[ + { + "text": "", + "ids": [] + }, + { + "text": "The capital of France is Paris.", + "ids": [ + 785, + 6722, + 315, + 9625, + 374, + 12095, + 13 + ] + }, + { + "text": "Hello, world!", + "ids": [ + 9707, + 11, + 1879, + 0 + ] + }, + { + "text": "don't I'll we've they're it's can't", + "ids": [ + 15007, + 944, + 358, + 3278, + 582, + 3003, + 807, + 2299, + 432, + 594, + 646, + 944 + ] + }, + { + "text": "DON'T SHOUT", + "ids": [ + 84641, + 17323, + 6434, + 3656 + ] + }, + { + "text": "spaces here and more", + "ids": [ + 44285, + 262, + 1588, + 257, + 323, + 981, + 803 + ] + }, + { + "text": " leading and trailing ", + "ids": [ + 220, + 6388, + 323, + 27748, + 256 + ] + }, + { + "text": "tabs\tand\tmore\ttabs", + "ids": [ + 30993, + 52477, + 2109, + 460, + 3244, + 3435 + ] + }, + { + "text": "newlines\n\nand\r\nwindows\r\nendings", + "ids": [ + 931, + 7969, + 271, + 437, + 319, + 27077, + 319, + 408, + 819 + ] + }, + { + "text": "mixed \n \t whitespace \n\n", + "ids": [ + 56685, + 715, + 19271, + 36372, + 4710 + ] + }, + { + "text": "1 12 123 1234 100000 3.14159", + "ids": [ + 16, + 220, + 16, + 17, + 220, + 16, + 17, + 18, + 220, + 16, + 17, + 18, + 19, + 220, + 16, + 15, + 15, + 15, + 15, + 15, + 220, + 18, + 13, + 16, + 19, + 16, + 20, + 24 + ] + }, + { + "text": "snake_case camelCase kebab-case", + "ids": [ + 72139, + 19096, + 49152, + 4207, + 1962, + 47722, + 38485 + ] + }, + { + "text": "for (int i = 0; i < n; ++i) { sum += a[i]; }", + "ids": [ + 1958, + 320, + 396, + 600, + 284, + 220, + 15, + 26, + 600, + 366, + 308, + 26, + 3443, + 72, + 8, + 314, + 2629, + 1421, + 264, + 989, + 5265, + 335 + ] + }, + { + "text": "café naïve résumé Zürich", + "ids": [ + 924, + 58858, + 94880, + 586, + 9333, + 1242, + 963, + 1863, + 5186, + 713 + ] + }, + { + "text": "Ünïcödé ßharp", + "ids": [ + 52491, + 77, + 37572, + 66, + 2956, + 128505, + 1683, + 253, + 71, + 7876 + ] + }, + { + "text": "你好世界,今天天气很好。", + "ids": [ + 108386, + 99489, + 3837, + 100644, + 104307, + 101243, + 1773 + ] + }, + { + "text": "こんにちは世界", + "ids": [ + 89015, + 99489 + ] + }, + { + "text": "Привет мир", + "ids": [ + 53645, + 26991, + 8178, + 137144 + ] + }, + { + "text": "emoji 😀 and 👨‍👩‍👧‍👦 family", + "ids": [ + 37523, + 90316, + 323, + 61804, + 101, + 378, + 235, + 145233, + 378, + 235, + 145665, + 378, + 235, + 145988, + 2997 + ] + }, + { + "text": "math ∑∫√≠≤ symbols", + "ids": [ + 10374, + 11995, + 239, + 145706, + 144336, + 145129, + 144570, + 17738 + ] + }, + { + "text": "<|begin_of_text|>hi<|eot_id|>", + "ids": [ + 27, + 91, + 7265, + 3575, + 4326, + 91, + 29, + 6023, + 27, + 91, + 68, + 354, + 842, + 91, + 29 + ] + }, + { + "text": "<|start_header_id|>user<|end_header_id|>\n\nWhat?<|eot_id|>", + "ids": [ + 27, + 91, + 2468, + 8757, + 842, + 91, + 29, + 872, + 27, + 91, + 408, + 8757, + 842, + 91, + 1339, + 3838, + 75414, + 91, + 68, + 354, + 842, + 91, + 29 + ] + }, + { + "text": "a<|eot_id|><|eot_id|>b", + "ids": [ + 64, + 27, + 91, + 68, + 354, + 842, + 91, + 1784, + 91, + 68, + 354, + 842, + 91, + 29, + 65 + ] + }, + { + "text": "URL: https://example.com/path?q=1&x=2#frag", + "ids": [ + 3144, + 25, + 3703, + 1110, + 8687, + 905, + 50976, + 43782, + 28, + 16, + 5, + 87, + 28, + 17, + 2, + 33198 + ] + }, + { + "text": "@user #hashtag $100 50% (parens) [brackets] {braces}", + "ids": [ + 31, + 872, + 671, + 4648, + 34311, + 400, + 16, + 15, + 15, + 220, + 20, + 15, + 4, + 320, + 3380, + 4412, + 8, + 508, + 1323, + 18382, + 60, + 314, + 1323, + 2434, + 92 + ] + } +] \ No newline at end of file diff --git a/reference/fixtures_qwen3_yarn/v0.npy b/reference/fixtures_qwen3_yarn/v0.npy new file mode 100644 index 0000000..7b0a376 Binary files /dev/null and b/reference/fixtures_qwen3_yarn/v0.npy differ diff --git a/src/capi/mlxforge.cpp b/src/capi/mlxforge.cpp index ed1884c..0c18a3f 100644 --- a/src/capi/mlxforge.cpp +++ b/src/capi/mlxforge.cpp @@ -179,6 +179,9 @@ mlxforge_engine* mlxforge_engine_create2(const char* model_spec, /* v9: multi-row GEMV decode kernels. Zero-init keeps the default (on). */ if (covered(&opts->skinny_mm + 1) && opts->skinny_mm != 0) cfg.skinny_mm = opts->skinny_mm > 0; + /* v10: RoPE-scaling JSON override. NULL/empty keeps the checkpoint's config. */ + if (covered(&opts->rope_scaling + 1) && opts->rope_scaling && *opts->rope_scaling) + cfg.rope_scaling = opts->rope_scaling; auto handle = std::make_unique(); handle->model_name = model_spec; diff --git a/src/capi/mlxforge.h b/src/capi/mlxforge.h index c0f16c1..c448b6e 100644 --- a/src/capi/mlxforge.h +++ b/src/capi/mlxforge.h @@ -50,8 +50,11 @@ extern "C" { * v8: mlxforge_engine_opts2.prefill_chunk (chunked-prefill interleaving, * default-on) — appended, struct_size-gated; no new symbols. * v9: mlxforge_engine_opts2.skinny_mm (multi-row GEMV decode kernels, - * default-on) — appended, struct_size-gated; no new symbols. */ -#define MLXFORGE_ABI_VERSION 9 + * default-on) — appended, struct_size-gated; no new symbols. + * v10: mlxforge_engine_opts2.rope_scaling (RoPE-scaling JSON override, + * yarn/linear long-context support) — appended, struct_size-gated; no + * new symbols. */ +#define MLXFORGE_ABI_VERSION 10 typedef struct mlxforge_engine mlxforge_engine; typedef struct mlxforge_request mlxforge_request; @@ -152,7 +155,17 @@ mlxforge_engine* mlxforge_engine_create(const char* model_spec, * MLX's tiled GEMM, which runs at a fraction of GEMV bandwidth there * (ml-explore/mlx#3661) — roughly 2x per-row decode throughput at small * batch sizes. On by default. 0 keeps the default; < 0 disables (stock - * matmul); > 0 enables. */ + * matmul); > 0 enables. + * + * rope_scaling (v10+) overrides the checkpoint's RoPE-scaling config with a + * JSON object (vLLM --rope-scaling shape), enabling long-context yarn/linear + * scaling on a stock checkpoint, e.g. + * {"rope_type":"yarn","factor":4.0,"original_max_position_embeddings":32768} + * (a yarn override may omit original_max_position_embeddings to scale the + * checkpoint's shipped context). NULL/empty = use the checkpoint's config. + * Unknown/unsupported rope_type values, hybrid (Qwen3.5) / vision-language + * models and GGUF checkpoints FAIL engine creation with a clear *err — there + * is never a silent fall-back to unscaled RoPE. */ typedef struct { size_t struct_size; /* caller sets sizeof(mlxforge_engine_opts2) */ int max_waiting; /* max queued requests; <= 0 => default (256) */ @@ -170,6 +183,9 @@ typedef struct { /* ---- v9 ---- */ int skinny_mm; /* multi-row GEMV decode kernels; 0 => default (on); < 0 => off (stock matmul); > 0 => on */ + /* ---- v10 ---- */ + const char* rope_scaling; /* RoPE-scaling JSON override; NULL/empty => use the + checkpoint's config */ } mlxforge_engine_opts2; /* Create an engine with extended options (v6+). Identical contract to diff --git a/src/core/config.cpp b/src/core/config.cpp index 6f26e0c..34007a9 100644 --- a/src/core/config.cpp +++ b/src/core/config.cpp @@ -47,21 +47,36 @@ std::vector parse_eos_ids(const nlohmann::json& j) { return {it->get()}; } -// Attempt to parse the optional "rope_scaling" sub-object, if present and is an object. -// Returns std::nullopt if absent or of incorrect type. -// Otherwise, fills out RopeScaling struct with available fields. -std::optional parse_rope_scaling(const nlohmann::json& j) { - auto it = j.find("rope_scaling"); - if (it == j.end() || !it->is_object()) return std::nullopt; +// Parse a rope-scaling JSON object into a RopeScaling. mlx_lm reads the type from +// "rope_type" or the legacy "type" key and treats a typeless object as "default"; +// mirror that so the two sides agree on what a config means. +RopeScaling rope_scaling_from_json(const nlohmann::json& o) { RopeScaling rs; - rs.rope_type = it->value("rope_type", std::string{}); - rs.factor = it->value("factor", 1.0f); - rs.high_freq_factor = it->value("high_freq_factor", 1.0f); - rs.low_freq_factor = it->value("low_freq_factor", 1.0f); - rs.original_max_position_embeddings = it->value("original_max_position_embeddings", 0); + rs.rope_type = o.value("rope_type", o.value("type", std::string{})); + if (rs.rope_type.empty()) rs.rope_type = "default"; + rs.factor = o.value("factor", 1.0f); + rs.high_freq_factor = o.value("high_freq_factor", 1.0f); + rs.low_freq_factor = o.value("low_freq_factor", 1.0f); + rs.original_max_position_embeddings = o.value("original_max_position_embeddings", 0); + rs.beta_fast = o.value("beta_fast", 32.0f); + rs.beta_slow = o.value("beta_slow", 1.0f); + rs.mscale = o.value("mscale", 1.0f); + rs.mscale_all_dim = o.value("mscale_all_dim", 0.0f); return rs; } +// Attempt to parse the optional rope-scaling sub-object. Lives under "rope_scaling" +// for most models, "rope_parameters" for Qwen3.5 — read both (like parse_mrope) so +// an unsupported type in either spelling is *seen* and rejected at validation, not +// silently dropped. Returns std::nullopt if neither is present. +std::optional parse_rope_scaling(const nlohmann::json& j) { + for (const char* key : {"rope_scaling", "rope_parameters"}) { + auto it = j.find(key); + if (it != j.end() && it->is_object()) return rope_scaling_from_json(*it); + } + return std::nullopt; +} + // RoPE base frequency. Most configs expose a top-level "rope_theta"; Qwen3.5 nests // it (with partial_rotary_factor) under a "rope_parameters" sub-object instead. // Prefer the sub-object when present, else fall back to the required top-level key. @@ -228,4 +243,60 @@ ModelConfig ModelConfig::from_file(const std::string& path) { return from_json(j); } +void validate_rope_scaling(const ModelConfig& cfg) { + if (!cfg.rope_scaling) return; + const RopeScaling& rs = *cfg.rope_scaling; + const std::string& t = rs.rope_type; + if (t.empty() || t == "default" || t == "llama3") return; + + if (t == "yarn" || t == "linear") { + // Supported only on the shared full-rotary text path (DecoderModel's + // precomputed-freqs RoPE). The other position schemes have their own rope + // code that would silently ignore the scaling — reject instead. + if (cfg.full_attention_interval > 0 || cfg.partial_rotary_factor != 1.0f) { + throw std::runtime_error("rope_scaling '" + t + + "' is not supported for hybrid / partial-rotary models"); + } + if (cfg.has_vision_tower() || !cfg.mrope_section.empty()) { + throw std::runtime_error("rope_scaling '" + t + + "' is not supported for M-RoPE vision models"); + } + if (cfg.rope_freq_factors) { + throw std::runtime_error("rope_scaling '" + t + + "' conflicts with checkpoint-baked rope frequency factors"); + } + if (rs.factor <= 0.0f) { + throw std::runtime_error("rope_scaling '" + t + "' requires factor > 0"); + } + if (t == "yarn" && rs.original_max_position_embeddings <= 0) { + throw std::runtime_error( + "rope_scaling 'yarn' requires original_max_position_embeddings > 0"); + } + return; + } + + throw std::runtime_error("unsupported rope_scaling rope_type '" + t + + "' (supported: default, llama3, yarn, linear)"); +} + +void apply_rope_scaling_override(ModelConfig& cfg, const std::string& json) { + nlohmann::json o; + try { + o = nlohmann::json::parse(json); + } catch (const nlohmann::json::exception& e) { + throw std::runtime_error(std::string("invalid rope_scaling override JSON: ") + e.what()); + } + if (!o.is_object()) { + throw std::runtime_error("rope_scaling override must be a JSON object, e.g. " + "{\"rope_type\":\"yarn\",\"factor\":4.0}"); + } + RopeScaling rs = rope_scaling_from_json(o); + // A yarn override without the original window means "scale the checkpoint's + // shipped context" — the natural reading for a stock (unscaled) model. + if (rs.rope_type == "yarn" && rs.original_max_position_embeddings <= 0) { + rs.original_max_position_embeddings = cfg.max_position_embeddings; + } + cfg.rope_scaling = rs; +} + } // namespace mlxforge diff --git a/src/core/config.h b/src/core/config.h index 5e97a1c..750d75e 100644 --- a/src/core/config.h +++ b/src/core/config.h @@ -24,11 +24,16 @@ struct QuantParams { /// Used by the RoPE stage to adjust rotational frequencies and handle /// position extrapolation for extended context. struct RopeScaling { - std::string rope_type; ///< Type of RoPE scaling ("llama3", etc.) + std::string rope_type; ///< Type of RoPE scaling ("llama3", "yarn", "linear", ...). float factor = 1.0f; ///< Primary scaling factor. - float high_freq_factor = 1.0f; ///< Scaling factor for high frequency components. - float low_freq_factor = 1.0f; ///< Scaling factor for low frequency components. + float high_freq_factor = 1.0f; ///< llama3: scaling factor for high frequency components. + float low_freq_factor = 1.0f; ///< llama3: scaling factor for low frequency components. int original_max_position_embeddings = 0; ///< Original context length before scaling. + // YaRN parameters (defaults mirror mlx_lm's YarnRoPE). + float beta_fast = 32.0f; ///< yarn: rotations bound for the interpolation ramp start. + float beta_slow = 1.0f; ///< yarn: rotations bound for the interpolation ramp end. + float mscale = 1.0f; ///< yarn: attention-scale numerator coefficient. + float mscale_all_dim = 0.0f; ///< yarn: attention-scale denominator coefficient. }; /// @brief Vision-tower (ViT) hyperparameters for a multimodal checkpoint. @@ -198,4 +203,20 @@ struct ModelConfig { static ModelConfig from_file(const std::string& path); }; +/// @brief Validate cfg.rope_scaling against the supported set. No scaling, "default" +/// and "llama3" always pass; "yarn"/"linear" pass only on the shared full-rotary +/// text path (rejected for hybrid/Qwen3.5, M-RoPE/vision and GGUF freq-factor +/// models); anything else throws. Replaces the old silent fall-through to +/// unscaled RoPE — an unsupported type must fail loading, never degrade. +/// @throws std::runtime_error naming the unsupported type or conflicting feature. +void validate_rope_scaling(const ModelConfig& cfg); + +/// @brief Replace cfg.rope_scaling with the parsed `json` object (vLLM --rope-scaling +/// semantics: full replacement, not a merge). For yarn overrides that omit +/// original_max_position_embeddings, defaults it to the checkpoint's +/// max_position_embeddings. Does not validate the type — call +/// validate_rope_scaling afterwards. +/// @throws std::runtime_error on malformed JSON or a non-object value. +void apply_rope_scaling_override(ModelConfig& cfg, const std::string& json); + } // namespace mlxforge diff --git a/src/model/decoder_model.cpp b/src/model/decoder_model.cpp index 48dada1..aac1de0 100644 --- a/src/model/decoder_model.cpp +++ b/src/model/decoder_model.cpp @@ -20,11 +20,26 @@ namespace { constexpr float kTwoPi = 6.283185307179586f; -// Precompute RoPE frequencies. For Llama-3.2 (rope_type "llama3") this applies -// the frequency rescaling that mlx_lm's Llama3RoPE precomputes and hands to -// fast::rope via `freqs` (with base disabled). Plain models fall back to the -// standard base**(2i/d) schedule. Returns head_dim/2 float32 values. -mx::array compute_rope_freqs(const ModelConfig& cfg) { +// YaRN attention-scale coefficient. Mirrors mlx_lm's yarn_get_mscale. +float yarn_get_mscale(float scale, float mscale) { + if (scale <= 1.0f) return 1.0f; + return 0.1f * mscale * std::log(scale) + 1.0f; +} + +} // namespace + +// Precompute the RoPE frequency schedule + YaRN mscale. For Llama-3.2 +// (rope_type "llama3") this applies the frequency rescaling that mlx_lm's +// Llama3RoPE precomputes and hands to fast::rope via `freqs` (with base +// disabled); "yarn" and "linear" mirror mlx_lm's YarnRoPE / nn.RoPE(scale= +// 1/factor) the same way. Plain models fall back to the standard base**(2i/d) +// schedule. freqs is head_dim/2 float32 values. +RopeSetup compute_rope_setup(const ModelConfig& cfg) { + // Defense-in-depth: the engine validates on the caller thread before model + // construction (a worker-thread throw would terminate); this catches direct + // construction paths (tests, CLI). + validate_rope_scaling(cfg); + const int hd = cfg.head_dim; const float base = cfg.rope_theta; @@ -42,12 +57,60 @@ mx::array compute_rope_freqs(const ModelConfig& cfg) { {static_cast(cfg.rope_freq_factors->size())}, mx::float32); freqs = mx::multiply(freqs, factors); mx::eval(freqs); - return freqs; + return {freqs, 1.0f}; } - if (!cfg.rope_scaling || cfg.rope_scaling->rope_type != "llama3") { + const std::string type = cfg.rope_scaling ? cfg.rope_scaling->rope_type : std::string{}; + + if (type == "linear") { + // mlx_lm uses nn.RoPE(scale=1/factor); fast::rope divides positions by + // freqs, so scaling positions by 1/factor equals scaling freqs by factor. + freqs = mx::multiply(freqs, mx::array(cfg.rope_scaling->factor)); mx::eval(freqs); - return freqs; + return {freqs, 1.0f}; + } + + if (type == "yarn") { + // Exact mirror of mlx_lm's YarnRoPE.__init__ (validated against the + // fixtures_qwen3_yarn golden freqs): blend interpolated (freqs * factor) + // and extrapolated (unscaled) frequencies with a linear ramp between the + // beta_fast/beta_slow correction dims. + const RopeScaling& rs = *cfg.rope_scaling; + const double orig = static_cast(rs.original_max_position_embeddings); + const double log_base = std::log(static_cast(base)); + auto correction_dim = [&](float num_rotations) { + return hd * std::log(orig / (num_rotations * static_cast(kTwoPi))) / + (2.0 * log_base); + }; + double low = std::max(std::floor(correction_dim(rs.beta_fast)), 0.0); + double high = std::min(std::ceil(correction_dim(rs.beta_slow)), static_cast(hd - 1)); + if (low == high) high += 0.001; // singularity guard (mlx_lm yarn_linear_ramp_mask) + + // freq_mask = 1 - clip((arange(hd/2) - low) / (high - low), 0, 1) + mx::array dim_idx = mx::arange(0, hd / 2, 1, mx::float32); + mx::array ramp = mx::clip( + mx::divide(mx::subtract(dim_idx, mx::array(static_cast(low))), + mx::array(static_cast(high - low))), + mx::array(0.0f), mx::array(1.0f)); + mx::array freq_mask = mx::subtract(mx::array(1.0f), ramp); + + // freqs = (inter * extra) / (inter * mask + extra * (1 - mask)), + // inter = factor * extra: unscaled where mask==1, interpolated where mask==0. + mx::array freq_extra = freqs; + mx::array freq_inter = mx::multiply(mx::array(rs.factor), freq_extra); + mx::array denom = mx::add(mx::multiply(freq_inter, freq_mask), + mx::multiply(freq_extra, mx::subtract(mx::array(1.0f), freq_mask))); + freqs = mx::divide(mx::multiply(freq_inter, freq_extra), denom); + mx::eval(freqs); + + const float mscale = + yarn_get_mscale(rs.factor, rs.mscale) / yarn_get_mscale(rs.factor, rs.mscale_all_dim); + return {freqs, mscale}; + } + + if (type != "llama3") { + mx::eval(freqs); + return {freqs, 1.0f}; } const RopeScaling& rs = *cfg.rope_scaling; @@ -80,11 +143,9 @@ mx::array compute_rope_freqs(const ModelConfig& cfg) { freqs = mx::where(is_medium, smooth_freqs, freqs); mx::eval(freqs); - return freqs; + return {freqs, 1.0f}; } -} // namespace - bool rope_array_offset_overload_available() { mx::array x = mx::zeros({1, 1, 1, 4}, mx::float16); mx::array offset = mx::array({0}, mx::int32); // per-row offset (B,) @@ -95,7 +156,7 @@ bool rope_array_offset_overload_available() { } DecoderModel::DecoderModel(ModelConfig config, Weights weights) - : cfg_(std::move(config)), w_(std::move(weights)), rope_freqs_(compute_rope_freqs(cfg_)) { + : cfg_(std::move(config)), w_(std::move(weights)), rope_(compute_rope_setup(cfg_)) { log::debug("DecoderModel: type={} layers={} hidden={} heads={}/{} head_dim={} vocab={} " "quantized={}", cfg_.model_type, cfg_.n_layers, cfg_.hidden, cfg_.n_heads, cfg_.n_kv_heads, @@ -146,14 +207,24 @@ mx::array DecoderModel::rms_norm(const mx::array& x, const mx::array& weight) co return mx::fast::rms_norm(x, std::optional(weight), cfg_.rms_eps); } +mx::array DecoderModel::mscale_input(const mx::array& x) const { + // YaRN multiplies Q/K by mscale *before* the rotation, mirroring mlx_lm's + // YarnRoPE (which scales the rope input). The rotation is linear, so this + // equals post-scaling, and it keeps q_rope0/k_rope0 fixture-exact. The + // scalar takes x's dtype so the result dtype is unchanged (weak-scalar + // semantics, like mlx_lm's python float). + if (rope_.mscale == 1.0f) return x; + return mx::multiply(x, mx::array(rope_.mscale, x.dtype())); +} + mx::array DecoderModel::apply_rope(const mx::array& x, int offset) const { - return mx::fast::rope(x, cfg_.head_dim, /*traditional=*/false, /*base=*/std::nullopt, - /*scale=*/1.0f, offset, rope_freqs_); + return mx::fast::rope(mscale_input(x), cfg_.head_dim, /*traditional=*/false, + /*base=*/std::nullopt, /*scale=*/1.0f, offset, rope_.freqs); } mx::array DecoderModel::apply_rope(const mx::array& x, const mx::array& offset) const { - return mx::fast::rope(x, cfg_.head_dim, /*traditional=*/false, /*base=*/std::nullopt, - /*scale=*/1.0f, offset, rope_freqs_); + return mx::fast::rope(mscale_input(x), cfg_.head_dim, /*traditional=*/false, + /*base=*/std::nullopt, /*scale=*/1.0f, offset, rope_.freqs); } mx::array DecoderModel::norm_qk_head(const mx::array& h, int /*layer*/, bool /*is_query*/) const { diff --git a/src/model/decoder_model.h b/src/model/decoder_model.h index d381ae4..8b38ddb 100644 --- a/src/model/decoder_model.h +++ b/src/model/decoder_model.h @@ -29,14 +29,31 @@ namespace mx = mlx::core; // true if it ran the overload successfully. bool rope_array_offset_overload_available(); +// Precomputed RoPE state: the per-dimension frequency schedule (head_dim/2 +// float32 values, fed to fast::rope via `freqs` with base disabled) plus the +// YaRN attention scale (1.0 for every other scheme). +struct RopeSetup { + mx::array freqs; + float mscale = 1.0f; +}; + +// Compute the RoPE frequency schedule + mscale for a config: GGUF baked factors, +// llama3 rescaling, yarn, linear, or the plain base**(2i/d) schedule. Validates +// cfg.rope_scaling and throws on unsupported types (see validate_rope_scaling). +// Exposed for the fixture-gated unit tests. +RopeSetup compute_rope_setup(const ModelConfig& cfg); + class DecoderModel { public: virtual ~DecoderModel() = default; const ModelConfig& config() const { return cfg_; } const Weights& weights() const { return w_; } - // Precomputed RoPE frequencies (llama3 rescaling), head_dim/2 float32 values. - const mx::array& rope_freqs() const { return rope_freqs_; } + // Precomputed RoPE frequencies (llama3/yarn/linear rescaling), head_dim/2 + // float32 values. + const mx::array& rope_freqs() const { return rope_.freqs; } + // YaRN attention scale, multiplied into Q/K before fast::rope (1.0 unless yarn). + float rope_mscale() const { return rope_.mscale; } // Embedding lookup: tokens (B, L) int32 -> (B, L, hidden) fp16. mx::array embed(const mx::array& tokens) const; @@ -121,6 +138,10 @@ class DecoderModel { // to route per layer. Same residual-stream shape (B, L, hidden) as the input. virtual mx::array feed_forward(const mx::array& x, int layer) const; + // YaRN pre-scales the RoPE input by rope_.mscale (see apply_rope); a no-op + // returning x unchanged for every other scheme (mscale == 1). + mx::array mscale_input(const mx::array& x) const; + // input RMSNorm -> Q/K/V projections -> reshape to heads (norm_qk_head on // Q/K), WITHOUT RoPE. QKV project_qkv(const mx::array& x, int layer) const; @@ -137,7 +158,7 @@ class DecoderModel { ModelConfig cfg_; Weights w_; - mx::array rope_freqs_; + RopeSetup rope_; bool skinny_mm_ = false; // see set_skinny_mm() }; diff --git a/src/runtime/engine.cpp b/src/runtime/engine.cpp index 32f3e39..75553c4 100644 --- a/src/runtime/engine.cpp +++ b/src/runtime/engine.cpp @@ -113,7 +113,7 @@ int validate_prefill_chunk(const EngineConfig& ec) { // For GGUF, this loads only the config/tokenizer fields from the GGUF header (not tensors). // For non-GGUF (MLX-style), loads config.json and tokenizer.json from disk. // This must be called on the main thread (MLX arrays are thread-locked to creating thread). -Engine::Loaded Engine::load_head(const std::string& spec) { +Engine::Loaded Engine::load_head(const std::string& spec, const std::string& rope_scaling) { Loaded out; // Resolve the physical directory for this model spec (could be local path or HF repo). @@ -144,6 +144,16 @@ Engine::Loaded Engine::load_head(const std::string& spec) { ); } + // RoPE-scaling override + validation, on the caller thread. The worker thread + // re-applies the same override when it rebuilds the config (make_factory) but + // cannot throw cleanly, so every rejection must happen here. + if (!rope_scaling.empty()) { + if (out.is_gguf) + throw std::runtime_error("rope_scaling override is not supported for GGUF models"); + apply_rope_scaling_override(out.config, rope_scaling); + } + validate_rope_scaling(out.config); + // Sniff embedding defaults from the on-disk sentence-transformers sidecar (no // such files inside a GGUF, so this is a no-op there). detect_embedding_defaults(out.dir, out.embed_pooling_default, out.embed_add_eos_default); @@ -155,8 +165,10 @@ Engine::Loaded Engine::load_head(const std::string& spec) { // MLX arrays must be created on the thread that will use them (the worker's). // - For GGUF: loads the weights from the GGUF file, then builds model. // - For non-GGUF: loads model config and weights, then builds model. -Worker::ModelFactory Engine::make_factory(std::string dir, bool is_gguf) { - return [dir = std::move(dir), is_gguf]() -> std::unique_ptr { +Worker::ModelFactory Engine::make_factory(std::string dir, bool is_gguf, + std::string rope_scaling) { + return [dir = std::move(dir), is_gguf, + rope_scaling = std::move(rope_scaling)]() -> std::unique_ptr { if (is_gguf) { // Parse config and load all weights from GGUF in one shot. GgufModel g = load_gguf_model(dir); @@ -164,6 +176,11 @@ Worker::ModelFactory Engine::make_factory(std::string dir, bool is_gguf) { } // Load config and weight tensors for legacy (non-GGUF) model format. ModelConfig wcfg = ModelConfig::from_file(dir + "/config.json"); + // Re-apply the rope-scaling override on the worker's copy of the config. + // load_head already applied + validated the identical config/override on the + // caller thread, so this cannot newly throw (a throw here would terminate — + // the worker run loop has no catch around the factory). + if (!rope_scaling.empty()) apply_rope_scaling_override(wcfg, rope_scaling); auto weights = load_weights(dir, wcfg); return create_model(std::move(wcfg), std::move(weights)); }; @@ -171,7 +188,7 @@ Worker::ModelFactory Engine::make_factory(std::string dir, bool is_gguf) { // Engine constructor: minimal form — loads config/tokenizer head from the given model spec. Engine::Engine(EngineConfig cfg) - : Engine(cfg, load_head(cfg.model_spec)) {} + : Engine(cfg, load_head(cfg.model_spec, cfg.rope_scaling)) {} // Engine constructor: explicit Loaded head (allows passing in preloaded config/tokenizer). // Sets up model name, config, tokenizer, and spawns worker thread for weights/model load. @@ -184,7 +201,8 @@ Engine::Engine(EngineConfig cfg, Loaded loaded) // Pass the tokenizer so the worker can build per-token byte strings for // constrained decoding. tok_ is initialized above and outlives worker_. // cfg_ is initialized above, so the KV-quant validation sees the model. - worker_(make_factory(std::move(loaded.dir), loaded.is_gguf), &scheduler_, &tok_, + worker_(make_factory(std::move(loaded.dir), loaded.is_gguf, cfg.rope_scaling), &scheduler_, + &tok_, validate_kv_quant(cfg, cfg_), validate_prefix_cache(cfg, cfg_, model_name_), validate_prefill_chunk(cfg), cfg.skinny_mm) { // Configure the max waiting requests for the batch scheduler. diff --git a/src/runtime/engine.h b/src/runtime/engine.h index 685d8fb..018e11b 100644 --- a/src/runtime/engine.h +++ b/src/runtime/engine.h @@ -57,6 +57,13 @@ struct EngineConfig { // On by default; logits may differ from the stock kernel at fp16-noise scale // (fp32 accumulation in a different order), token-equality gated in tests. bool skinny_mm = true; + // RoPE-scaling override (vLLM --rope-scaling style): a JSON object replacing + // the checkpoint's rope_scaling config, e.g. + // {"rope_type":"yarn","factor":4.0,"original_max_position_embeddings":32768} + // Empty = use the checkpoint's config as-is. Like kv_bits, unsupported setups + // (unknown rope_type, hybrid/vision models, GGUF) FAIL engine creation — + // never a silent fall-back to unscaled RoPE. + std::string rope_scaling; }; // Per-call embedding options. The two int fields are tri-state: -1 means "use @@ -123,10 +130,13 @@ class Engine { bool embed_add_eos_default = false; }; - // Resolve model spec path, parse config/tokenizer, etc. - static Loaded load_head(const std::string& spec); + // Resolve model spec path, parse config/tokenizer, etc. Applies the optional + // rope-scaling override and validates it — this is the caller-thread rejection + // point for unsupported rope configs (the worker thread cannot throw cleanly). + static Loaded load_head(const std::string& spec, const std::string& rope_scaling); // Factory builder: creates a Worker::ModelFactory, handling weight loading with proper backend - static Worker::ModelFactory make_factory(std::string dir, bool is_gguf); + static Worker::ModelFactory make_factory(std::string dir, bool is_gguf, + std::string rope_scaling); // Private delegating ctor, used internally after head-loading step is complete. Engine(EngineConfig cfg, Loaded loaded); diff --git a/src/server/config.cpp b/src/server/config.cpp index 9140ba3..30cd671 100644 --- a/src/server/config.cpp +++ b/src/server/config.cpp @@ -59,7 +59,7 @@ ServerConfig ServerConfig::from_file(const std::string& path) { static const std::set kKnownKeys = { "model", "host", "port", "max_ctx", "max_waiting", "kv_budget", "kv_bits", "prefix_cache", "kv_block", "kv_pool", "kv_spill_dir", "kv_spill_bytes", - "prefill_chunk", "skinny_mm"}; + "prefill_chunk", "skinny_mm", "rope_scaling"}; for (const auto& [key, _] : j.items()) { if (kKnownKeys.find(key) == kKnownKeys.end()) { throw std::runtime_error("config file: unknown key '" + key + "' in '" + path + "'"); @@ -115,6 +115,15 @@ ServerConfig ServerConfig::from_file(const std::string& path) { throw std::runtime_error("config file: 'prefill_chunk' must be >= 0 (0 = monolithic)"); } if (j.contains("skinny_mm")) c.skinny_mm = require_type(j, "skinny_mm"); + if (j.contains("rope_scaling")) { + // Accept either an inline JSON object (the natural config-file spelling) or + // a pre-serialized string; the engine parses + validates the contents. + const nlohmann::json& rs = j.at("rope_scaling"); + if (rs.is_object()) + c.rope_scaling = rs.dump(); + else + c.rope_scaling = require_type(j, "rope_scaling"); + } return c; } @@ -167,6 +176,7 @@ ServerConfig ServerConfig::parse(const std::vector& args) { env_long("MLXFORGE_KV_SPILL_BYTES", static_cast(c.kv_spill_bytes))); c.prefill_chunk = static_cast(env_long("MLXFORGE_PREFILL_CHUNK", c.prefill_chunk)); c.skinny_mm = env_long("MLXFORGE_SKINNY_MM", c.skinny_mm ? 1 : 0) != 0; + c.rope_scaling = env_or("MLXFORGE_ROPE_SCALING", c.rope_scaling); // Helper: extract value for a flag (accepts "--flag value" or "--flag=value") auto value_of = [&](const std::string& a, size_t& i) -> std::string { @@ -213,6 +223,8 @@ ServerConfig ServerConfig::parse(const std::vector& args) { c.prefill_chunk = std::stoi(value_of(a, i)); else if (flag == "--skinny-mm") c.skinny_mm = std::stoi(value_of(a, i)) != 0; + else if (flag == "--rope-scaling") + c.rope_scaling = value_of(a, i); else throw std::runtime_error("unknown flag: " + flag); } diff --git a/src/server/config.h b/src/server/config.h index face2d4..ccab3d3 100644 --- a/src/server/config.h +++ b/src/server/config.h @@ -51,6 +51,10 @@ struct ServerConfig { // Multi-row GEMV decode kernels for small batched-decode matmuls (default on). bool skinny_mm = true; + // RoPE-scaling JSON override (vLLM --rope-scaling style), passed through to + // the engine. Empty = use the checkpoint's rope_scaling config. + std::string rope_scaling; + // Parses command line arguments (the model via -m/--model, plus optional flags as --flag value or --flag=value), // layering configuration sources by precedence (lowest to highest): // struct defaults < config file (-c/--config) < environment variables < CLI flags. @@ -58,7 +62,7 @@ struct ServerConfig { // 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. // Throws std::runtime_error if an unknown or malformed flag is encountered. static ServerConfig parse(const std::vector& args); @@ -66,7 +70,7 @@ struct ServerConfig { // with struct defaults filling any keys the file omits. Recognized keys // (snake_case): "model", "host", "port", "max_ctx", "max_waiting", "kv_budget", // "kv_bits", "prefix_cache", "kv_block", "kv_pool", "kv_spill_dir", - // "kv_spill_bytes", "prefill_chunk", "skinny_mm". + // "kv_spill_bytes", "prefill_chunk", "skinny_mm", "rope_scaling". // Validates before applying: rejects unknown keys, wrong types, and out-of-range // values. Throws std::runtime_error (with the file path / offending key) on any // failure to open, parse, or validate. diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 47843e2..4278430 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -40,6 +40,8 @@ add_executable(mlxforge_tests runtime/bucketing_test.cpp scheduler/validation_test.cpp model/qwen3_test.cpp + model/rope_scaling_test.cpp + model/qwen3_yarn_test.cpp model/qwen3_moe_test.cpp model/qwen3_5_test.cpp model/vit_test.cpp @@ -159,6 +161,7 @@ target_compile_definitions(mlxforge_tests PRIVATE MLXFORGE_TEST_FIXTURES_DIR="${CMAKE_CURRENT_SOURCE_DIR}/fixtures" MLXFORGE_REF_FIXTURES_DIR="${CMAKE_SOURCE_DIR}/reference/fixtures" MLXFORGE_REF_FIXTURES_DIR_QWEN3="${CMAKE_SOURCE_DIR}/reference/fixtures_qwen3" + MLXFORGE_REF_FIXTURES_DIR_QWEN3_YARN="${CMAKE_SOURCE_DIR}/reference/fixtures_qwen3_yarn" MLXFORGE_REF_FIXTURES_DIR_QWEN3_MOE="${CMAKE_SOURCE_DIR}/reference/fixtures_qwen3_moe" MLXFORGE_REF_FIXTURES_DIR_QWEN3_5="${CMAKE_SOURCE_DIR}/reference/fixtures_qwen3_5" MLXFORGE_REF_FIXTURES_DIR_QWEN3_VL="${CMAKE_SOURCE_DIR}/reference/fixtures_qwen3_vl" diff --git a/tests/capi/capi_test.cpp b/tests/capi/capi_test.cpp index c3e529b..959f1f9 100644 --- a/tests/capi/capi_test.cpp +++ b/tests/capi/capi_test.cpp @@ -425,3 +425,45 @@ TEST_CASE("C ABI v9 skinny_mm: kernel-on and stock-matmul engines agree") { CHECK(run_with(1) == stock); // explicit on CHECK(run_with(0) == stock); // the default (on) } + +TEST_CASE("C ABI v10 rope_scaling: yarn override loads; unknown types fail creation") { + if (!model_available()) { + MESSAGE("MLXFORGE_MODEL_DIR not present; skipping"); + return; + } + // An unsupported rope_type must fail creation with a clear *err on the caller + // thread — never a silent fall-back to unscaled RoPE. + char* err = nullptr; + mlxforge_engine_opts2 opts = {}; + opts.struct_size = sizeof(opts); + opts.rope_scaling = "{\"rope_type\":\"dynamic\",\"factor\":2.0}"; + mlxforge_engine* eng = mlxforge_engine_create2(model_dir().c_str(), &opts, &err); + CHECK(eng == nullptr); + REQUIRE(err != nullptr); + CHECK(std::string(err).find("unsupported rope_scaling") != std::string::npos); + mlxforge_string_free(err); + + // Malformed JSON is rejected the same way. + err = nullptr; + opts.rope_scaling = "{not json"; + eng = mlxforge_engine_create2(model_dir().c_str(), &opts, &err); + CHECK(eng == nullptr); + REQUIRE(err != nullptr); + CHECK(std::string(err).find("rope_scaling") != std::string::npos); + mlxforge_string_free(err); + + // A valid yarn override creates an engine that decodes coherently. Llama-3.2 + // ships llama3 scaling, so the override also exercises full replacement. + err = nullptr; + opts.rope_scaling = + "{\"rope_type\":\"yarn\",\"factor\":4.0,\"original_max_position_embeddings\":8192}"; + eng = mlxforge_engine_create2(model_dir().c_str(), &opts, &err); + REQUIRE_MESSAGE(eng != nullptr, (err ? err : "engine_create2 failed")); + mlxforge_sampling s = {}; + s.max_tokens = 4; + mlxforge_request* r = mlxforge_submit_text(eng, "Hello", &s, &err); + REQUIRE_MESSAGE(r != nullptr, (err ? err : "submit failed")); + CHECK(drain(r).size() > 0); + mlxforge_request_free(r); + mlxforge_engine_free(eng); +} diff --git a/tests/model/qwen3_yarn_test.cpp b/tests/model/qwen3_yarn_test.cpp new file mode 100644 index 0000000..cdf24d8 --- /dev/null +++ b/tests/model/qwen3_yarn_test.cpp @@ -0,0 +1,89 @@ +// Qwen3 + yarn golden-reference checks: the same Qwen3-0.6B weights with the +// yarn rope_scaling from the fixtures_qwen3_yarn manifest injected through the +// engine's override path (apply_rope_scaling_override), gated end to end against +// mlx-lm loaded with the identical model_config injection — the yarn-rescaled +// attention front-half (freqs + mscale-on-input), the full-forward argmax, and +// the greedy token stream. Self-skips unless both the Qwen3 model +// (MLXFORGE_MODEL_DIR_QWEN3) and the committed yarn fixtures are present. +#include + +#include +#include + +#include "mlx/ops.h" +#include "runtime/single_stream.h" +#include "support/model_fixture.h" +#include "support/reference.h" + +using namespace mlxforge::test; + +namespace { +bool yarn_fixtures_present() { + return std::ifstream(qwen3_yarn_ref_path("manifest.json")).good(); +} +bool yarn_ready() { return qwen3_model_available() && yarn_fixtures_present(); } +} // namespace + +TEST_CASE("Qwen3+yarn: rescaled attention front-half matches the reference") { + if (!yarn_ready()) { + MESSAGE("Qwen3 model / yarn fixtures not present; skipping golden-reference check"); + return; + } + mlxforge::Qwen3Model& model = shared_qwen3_yarn_model(); + + // The override took: yarn freqs schedule + attention mscale, both vs mlx-lm. + assert_close(model.rope_freqs(), load_qwen3_yarn_npy("rope_freqs.npy"), /*rtol=*/1e-5f, + /*atol=*/1e-3f); + mx::array ref_mscale = load_qwen3_yarn_npy("rope_mscale.npy"); + mx::eval(ref_mscale); + CHECK(model.rope_mscale() == doctest::Approx(ref_mscale.item()).epsilon(1e-6)); + + std::vector ids = load_qwen3_yarn_token_ids("prompt_0_ids.npy"); + mx::array tokens(ids.data(), {1, static_cast(ids.size())}, mx::int32); + mx::array emb = model.embed(tokens); + assert_close(emb, load_qwen3_yarn_npy("embeddings.npy")); + + // Post-(q_norm + mscale + yarn-RoPE) Q/K. If the mscale were dropped or the + // freqs unscaled, q/k diverge from the reference here. + mlxforge::DecoderModel::QKV qkv = model.attn_qkv(emb, /*layer=*/0); + assert_close(qkv.q, load_qwen3_yarn_npy("q_rope0.npy")); + assert_close(qkv.k, load_qwen3_yarn_npy("k_rope0.npy")); + assert_close(qkv.v, load_qwen3_yarn_npy("v0.npy")); +} + +TEST_CASE("Qwen3+yarn: full forward logits + first-token argmax match the reference") { + if (!yarn_ready()) { + MESSAGE("Qwen3 model / yarn fixtures not present; skipping golden-reference check"); + return; + } + mlxforge::Qwen3Model& model = shared_qwen3_yarn_model(); + std::vector ids = load_qwen3_yarn_token_ids("prompt_0_ids.npy"); + const int T = static_cast(ids.size()); + mx::array tokens(ids.data(), {1, T}, mx::int32); + + mx::array logits = model.forward(tokens); // (1, T, vocab) + const int vocab = logits.shape()[2]; + mx::array last = mx::reshape(mx::slice(logits, {0, T - 1, 0}, {1, T, vocab}), {1, vocab}); + // Same loosened bound as the plain Qwen3 gate (28 layers of fp16 drift, now + // also mscale^2-scaled); the exact argmax below is the real correctness gate. + assert_close(last, load_qwen3_yarn_npy("logits_last.npy"), /*rtol=*/3e-2f, /*atol=*/3e-2f); + + std::vector argmax = load_qwen3_yarn_token_ids("argmax.npy"); + mx::array got = mx::astype(mx::argmax(last, /*axis=*/-1), mx::int32); + mx::eval(got); + CHECK(got.data()[0] == argmax[0]); +} + +TEST_CASE("Qwen3+yarn: greedy stream matches mlx-lm token-for-token") { + if (!yarn_ready()) { + MESSAGE("Qwen3 model / yarn fixtures not present; skipping golden-reference check"); + return; + } + mlxforge::Qwen3Model& model = shared_qwen3_yarn_model(); + std::vector prompt = load_qwen3_yarn_token_ids("prompt_0_ids.npy"); + std::vector ref = load_qwen3_yarn_token_ids("greedy_tokens.npy"); + + mlxforge::GenerateResult r = mlxforge::greedy_generate( + model, prompt, /*max_tokens=*/static_cast(ref.size()), model.config().eos_token_ids); + assert_tokens_equal(r.tokens, ref); +} diff --git a/tests/model/rope_scaling_test.cpp b/tests/model/rope_scaling_test.cpp new file mode 100644 index 0000000..9b7998a --- /dev/null +++ b/tests/model/rope_scaling_test.cpp @@ -0,0 +1,221 @@ +// RoPE-scaling unit gates (no model download needed): the yarn/linear frequency +// schedules and the yarn mscale vs the mlx-lm fixtures (fixture-to-fixture — the +// committed q_pre0 tensor pushed through our freqs must reproduce the committed +// q_rope0), plus the validation/override behavior: unknown rope_type values and +// unsupported model shapes must throw, never silently fall back to unscaled RoPE. +#include + +#include +#include +#include +#include + +#include "core/config.h" +#include "mlx/fast.h" +#include "mlx/ops.h" +#include "model/decoder_model.h" +#include "support/reference.h" + +using namespace mlxforge::test; + +namespace { + +bool yarn_fixtures_present() { + return std::ifstream(qwen3_yarn_ref_path("manifest.json")).good(); +} + +// Minimal config in the Qwen3-0.6B rope shape (the only fields +// compute_rope_setup reads besides rope_scaling). +mlxforge::ModelConfig base_cfg() { + mlxforge::ModelConfig c; + c.head_dim = 128; + c.rope_theta = 1000000.0f; + c.max_position_embeddings = 40960; + return c; +} + +// The injected yarn recipe from dump_ref.py's qwen3_yarn spec. +mlxforge::RopeScaling yarn_scaling() { + mlxforge::RopeScaling rs; + rs.rope_type = "yarn"; + rs.factor = 4.0f; + rs.original_max_position_embeddings = 32768; + return rs; +} + +// Plain base**(2i/d) schedule, the unscaled reference for the no-op cases. +mx::array plain_freqs(const mlxforge::ModelConfig& c) { + mx::array idx = mx::arange(0, c.head_dim, 2, mx::float32); + mx::array freqs = mx::power(mx::array(c.rope_theta), + mx::divide(idx, mx::array(static_cast(c.head_dim)))); + mx::eval(freqs); + return freqs; +} + +mx::array rope_with_freqs(const mx::array& x, const mx::array& freqs) { + return mx::fast::rope(x, x.shape().back(), /*traditional=*/false, /*base=*/std::nullopt, + /*scale=*/1.0f, /*offset=*/0, freqs); +} + +} // namespace + +TEST_CASE("yarn freqs + mscale match mlx-lm's YarnRoPE") { + if (!yarn_fixtures_present()) { + MESSAGE("fixtures_qwen3_yarn not present; skipping"); + return; + } + mlxforge::ModelConfig c = base_cfg(); + c.rope_scaling = yarn_scaling(); + mlxforge::RopeSetup setup = mlxforge::compute_rope_setup(c); + + // Both sides compute the schedule in float32 from the same formula, so the + // tolerance is far tighter than the fp16 tensor gates. + assert_close(setup.freqs, load_qwen3_yarn_npy("rope_freqs.npy"), /*rtol=*/1e-5f, + /*atol=*/1e-3f); + + mx::array ref_mscale = load_qwen3_yarn_npy("rope_mscale.npy"); + mx::eval(ref_mscale); + CHECK(setup.mscale == doctest::Approx(ref_mscale.item()).epsilon(1e-6)); + // Qwen3 yarn defaults (mscale=1, mscale_all_dim=0) reduce to 0.1*ln(factor)+1. + CHECK(setup.mscale == doctest::Approx(0.1f * std::log(4.0f) + 1.0f).epsilon(1e-6)); +} + +TEST_CASE("yarn rope application reproduces the reference q_rope0/k_rope0") { + if (!yarn_fixtures_present()) { + MESSAGE("fixtures_qwen3_yarn not present; skipping"); + return; + } + mlxforge::ModelConfig c = base_cfg(); + c.rope_scaling = yarn_scaling(); + mlxforge::RopeSetup setup = mlxforge::compute_rope_setup(c); + + // Fixture-to-fixture: mscale * q_pre0 through fast::rope with our freqs must + // equal mlx-lm's attn.rope(q) (YarnRoPE scales its input by mscale). + mx::array q = load_qwen3_yarn_npy("q_pre0.npy"); + mx::array qs = mx::multiply(q, mx::array(setup.mscale, q.dtype())); + assert_close(rope_with_freqs(qs, setup.freqs), load_qwen3_yarn_npy("q_rope0.npy")); +} + +TEST_CASE("linear rope matches mlx-lm's nn.RoPE(scale=1/factor)") { + if (!yarn_fixtures_present()) { + MESSAGE("fixtures_qwen3_yarn not present; skipping"); + return; + } + mlxforge::ModelConfig c = base_cfg(); + mlxforge::RopeScaling rs; + rs.rope_type = "linear"; + rs.factor = 4.0f; // the factor dump_ref.py's linear oracle uses + c.rope_scaling = rs; + mlxforge::RopeSetup setup = mlxforge::compute_rope_setup(c); + CHECK(setup.mscale == 1.0f); + + mx::array q = load_qwen3_yarn_npy("q_pre0.npy"); + assert_close(rope_with_freqs(q, setup.freqs), load_qwen3_yarn_npy("q_rope0_linear.npy")); +} + +TEST_CASE("default / missing rope_scaling keep the plain schedule") { + mlxforge::ModelConfig c = base_cfg(); + mlxforge::RopeSetup none = mlxforge::compute_rope_setup(c); + CHECK(none.mscale == 1.0f); + assert_close(none.freqs, plain_freqs(c), /*rtol=*/1e-6f, /*atol=*/1e-3f); + + mlxforge::RopeScaling rs; + rs.rope_type = "default"; + c.rope_scaling = rs; + mlxforge::RopeSetup dflt = mlxforge::compute_rope_setup(c); + CHECK(dflt.mscale == 1.0f); + assert_close(dflt.freqs, plain_freqs(c), /*rtol=*/1e-6f, /*atol=*/1e-3f); +} + +TEST_CASE("unknown rope_type values are rejected, never silently unscaled") { + mlxforge::ModelConfig c = base_cfg(); + mlxforge::RopeScaling rs; + rs.factor = 2.0f; + rs.original_max_position_embeddings = 4096; + + for (const char* t : {"dynamic", "longrope", "su"}) { + rs.rope_type = t; + c.rope_scaling = rs; + CHECK_THROWS_WITH_AS(mlxforge::validate_rope_scaling(c), + ("unsupported rope_scaling rope_type '" + std::string(t) + + "' (supported: default, llama3, yarn, linear)") + .c_str(), + std::runtime_error); + // compute_rope_setup re-validates (defense for direct construction paths). + CHECK_THROWS_AS(mlxforge::compute_rope_setup(c), std::runtime_error); + } +} + +TEST_CASE("yarn/linear are rejected off the shared full-rotary text path") { + mlxforge::RopeScaling rs = yarn_scaling(); + + // Hybrid (Qwen3.5): own partial rope, never reads the shared freqs. + mlxforge::ModelConfig hybrid = base_cfg(); + hybrid.full_attention_interval = 4; + hybrid.rope_scaling = rs; + CHECK_THROWS_AS(mlxforge::validate_rope_scaling(hybrid), std::runtime_error); + + mlxforge::ModelConfig partial = base_cfg(); + partial.partial_rotary_factor = 0.5f; + partial.rope_scaling = rs; + CHECK_THROWS_AS(mlxforge::validate_rope_scaling(partial), std::runtime_error); + + // Vision / M-RoPE (Qwen3-VL): 3D positions, hand-rolled rotation. + mlxforge::ModelConfig vision = base_cfg(); + vision.vision = mlxforge::VisionConfig{}; + vision.rope_scaling = rs; + CHECK_THROWS_AS(mlxforge::validate_rope_scaling(vision), std::runtime_error); + + mlxforge::ModelConfig mrope = base_cfg(); + mrope.mrope_section = {24, 20, 20}; + mrope.rope_scaling = rs; + CHECK_THROWS_AS(mlxforge::validate_rope_scaling(mrope), std::runtime_error); + + // GGUF checkpoints bake llama3 factors into rope_freqs.weight — conflicting. + mlxforge::ModelConfig gguf = base_cfg(); + gguf.rope_freq_factors = std::vector(64, 1.0f); + gguf.rope_scaling = rs; + CHECK_THROWS_AS(mlxforge::validate_rope_scaling(gguf), std::runtime_error); + + // Required parameters. + mlxforge::ModelConfig no_orig = base_cfg(); + mlxforge::RopeScaling bad = rs; + bad.original_max_position_embeddings = 0; + no_orig.rope_scaling = bad; + CHECK_THROWS_AS(mlxforge::validate_rope_scaling(no_orig), std::runtime_error); + + mlxforge::ModelConfig bad_factor = base_cfg(); + bad = rs; + bad.factor = 0.0f; + bad_factor.rope_scaling = bad; + CHECK_THROWS_AS(mlxforge::validate_rope_scaling(bad_factor), std::runtime_error); +} + +TEST_CASE("apply_rope_scaling_override: replacement semantics and defaults") { + mlxforge::ModelConfig c = base_cfg(); + + // Malformed JSON / non-object values are clear errors. + CHECK_THROWS_AS(mlxforge::apply_rope_scaling_override(c, "{not json"), std::runtime_error); + CHECK_THROWS_AS(mlxforge::apply_rope_scaling_override(c, "42"), std::runtime_error); + + // A yarn override without the original window scales the checkpoint's + // shipped context (max_position_embeddings). + mlxforge::apply_rope_scaling_override(c, R"({"rope_type":"yarn","factor":4.0})"); + REQUIRE(c.rope_scaling.has_value()); + CHECK(c.rope_scaling->rope_type == "yarn"); + CHECK(c.rope_scaling->factor == 4.0f); + CHECK(c.rope_scaling->original_max_position_embeddings == 40960); + CHECK(c.rope_scaling->beta_fast == 32.0f); // mlx-lm YarnRoPE defaults + CHECK(c.rope_scaling->beta_slow == 1.0f); + mlxforge::validate_rope_scaling(c); // the resulting config is valid + + // Full replacement (vLLM semantics), and the legacy "type" key is accepted. + mlxforge::apply_rope_scaling_override(c, R"({"type":"linear","factor":2.0})"); + CHECK(c.rope_scaling->rope_type == "linear"); + CHECK(c.rope_scaling->factor == 2.0f); + + // An unknown type in an override survives parsing but fails validation — + // the engine calls validate right after. + mlxforge::apply_rope_scaling_override(c, R"({"rope_type":"dynamic","factor":2.0})"); + CHECK_THROWS_AS(mlxforge::validate_rope_scaling(c), std::runtime_error); +} diff --git a/tests/support/model_fixture.h b/tests/support/model_fixture.h index 4b9fd78..5763757 100644 --- a/tests/support/model_fixture.h +++ b/tests/support/model_fixture.h @@ -8,6 +8,8 @@ #include #include +#include + #include "core/config.h" #include "core/weights.h" #include "model/llama.h" @@ -53,6 +55,28 @@ inline Qwen3Model& shared_qwen3_model() { return model; } +// The Qwen3 weights again, with the yarn rope_scaling recorded in the +// fixtures_qwen3_yarn manifest injected through the engine's override path — +// the C++ mirror of dump_ref.py's model_config injection (no checkpoint on the +// Hub ships yarn, so both sides inject the identical object). +inline std::string qwen3_yarn_rope_json() { + std::ifstream f(std::string(MLXFORGE_REF_FIXTURES_DIR_QWEN3_YARN) + "/manifest.json"); + nlohmann::json m; + f >> m; + return m.at("model_config").at("rope_scaling").dump(); +} + +inline Qwen3Model& shared_qwen3_yarn_model() { + static Qwen3Model model = [] { + ModelConfig cfg = ModelConfig::from_file(qwen3_model_dir() + "/config.json"); + apply_rope_scaling_override(cfg, qwen3_yarn_rope_json()); + validate_rope_scaling(cfg); + Weights w = load_weights(qwen3_model_dir(), cfg); + return Qwen3Model(std::move(cfg), std::move(w)); + }(); + return model; +} + // Same, for the Qwen3 MoE model (sparse expert-routing integration tests). inline std::string qwen3_moe_model_dir() { return MLXFORGE_MODEL_DIR_QWEN3_MOE; } diff --git a/tests/support/reference.h b/tests/support/reference.h index 9c81f99..d225530 100644 --- a/tests/support/reference.h +++ b/tests/support/reference.h @@ -52,6 +52,18 @@ inline std::vector load_qwen3_token_ids(const std::string& name) { return load_token_ids_at(MLXFORGE_REF_FIXTURES_DIR_QWEN3, name); } +// Same accessors against the Qwen3 yarn fixture set (reference/fixtures_qwen3_yarn): +// the same Qwen3 weights with an injected yarn rope_scaling (see dump_ref.py). +inline std::string qwen3_yarn_ref_path(const std::string& name) { + return std::string(MLXFORGE_REF_FIXTURES_DIR_QWEN3_YARN) + "/" + name; +} +inline mx::array load_qwen3_yarn_npy(const std::string& name) { + return mx::load(qwen3_yarn_ref_path(name)); +} +inline std::vector load_qwen3_yarn_token_ids(const std::string& name) { + return load_token_ids_at(MLXFORGE_REF_FIXTURES_DIR_QWEN3_YARN, name); +} + // Same accessors against the Qwen3 MoE fixture set (reference/fixtures_qwen3_moe). inline std::string qwen3_moe_ref_path(const std::string& name) { return std::string(MLXFORGE_REF_FIXTURES_DIR_QWEN3_MOE) + "/" + name;