From f0364d12b1fecf1bfca3f06d7e65ec750ff99622 Mon Sep 17 00:00:00 2001 From: fab2s Date: Thu, 6 Aug 2026 00:08:30 +0200 Subject: [PATCH 1/4] Apply rotary to MLA's rope dims on models that are not NoPE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `src/` implemented no rotary: every occurrence of `rope` was `qk_rope` used as a width, and `rope_theta`, `rope_scaling` and `mla_use_nope` were read nowhere. That is correct for the Kimi models, which set `mla_use_nope` and pass those dims through unrotated, and wrong for a DeepSeek-V3 checkpoint, which sets no such flag and ships `rope_theta` with YaRN scaling. In MLA those dims are the only positional signal — the nope dims are position-free by construction — so a V3-family container attended over an unordered sequence. It is quiet: lexically determined answers still come out right, so a single-turn factual prompt looks correct. Add a second turn boundary and the model emits an empty assistant turn, `<|im_end|>` at p=0.968 on Kimi-K2, because it cannot order the prompt. `rope_init` builds `inv_freq` with YaRN's ramp and takes `mscale_all_dim` squared onto the attention scale, following `DeepseekV3YarnRotaryEmbedding`. Two details do not survive paraphrase: the rotation is GPT-J interleaved, pairing `x[2j]` with `x[2j+1]` — upstream reaches the same arithmetic by de-interleaving before a half-split rotate, so applying the half-split form directly pairs the wrong dims and still yields finite, weight-shaped output — and YaRN rescales `inv_freq` globally, so it applies from position 0 rather than only at long context. The k-side slice is rotated before it enters the latent cache, because a cached entry is reused by every later query and carries its own token's position. The absorbed `kv_b_proj` identity touches only the nope half and is unaffected. Rotation is skipped entirely when `mla_use_nope` is set, so Kimi-Linear and K3 are untouched by construction. A container needing rotation on a slice wider than `2 * WASTE_MAX_ROPE_HALF` is refused at load rather than run unrotated. `tools/kimi_ref.py` cannot serve a V3 config — it indexes `linear_attn_config` unconditionally and applies no rotary — so `tools/deepseek_ref.py` is the oracle for this path, on the same contract: weights read from the container, so quantization error cancels and a diff measures arithmetic. With `--no-rope --no-mscale` it reproduces the engine's pre-fix layer-0 residual to 0.0001% rel L2, and its full-depth top-1 to p=0.967985 against 0.968. --- src/model.c | 112 +++++++++++++- src/model.h | 16 ++ tools/deepseek_ref.py | 331 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 456 insertions(+), 3 deletions(-) create mode 100644 tools/deepseek_ref.py diff --git a/src/model.c b/src/model.c index 51ef35b..7d8928f 100644 --- a/src/model.c +++ b/src/model.c @@ -848,6 +848,58 @@ static int cfg_sane(const waste_config *c) return 1; } +/* inv_freq for the rope dims, plus YaRN's factor on the attention scale. + * + * Both follow DeepseekV3YarnRotaryEmbedding. Two details do not survive + * paraphrase: + * - mscale appears twice with different meanings. cos/sin carry + * mscale / mscale_all_dim, which is 1 whenever the two are equal (K2 sets + * both to 1), so they are left alone here. The attention scale carries + * mscale_all_dim SQUARED, which is 1.8133x on K2. + * - YaRN rescales inv_freq globally, so it applies from position 0. It is + * not a long-context-only correction that a short prompt can ignore. + */ +static void rope_init(waste_config *c, const js_doc *d, int cfg) +{ + const double PI = 3.14159265358979323846; + c->att_mul = 1.0f; + c->mla_nope = js_get(d, cfg, "mla_use_nope") >= 0; + const int dim = c->qk_rope, half = dim / 2; + if (c->mla_nope || half <= 0 || half > WASTE_MAX_ROPE_HALF) return; + + const double base = js_num(d, js_get(d, cfg, "rope_theta"), 10000.0); + for (int j = 0; j < half; j++) + c->rope_inv_freq[j] = (float)(1.0 / pow(base, (double)(2 * j) / dim)); + + const int rs = js_get(d, cfg, "rope_scaling"); + if (rs < 0) return; + char type[16]; + js_str(d, js_get(d, rs, "type"), type, sizeof type); + const double factor = js_num(d, js_get(d, rs, "factor"), 1.0); + if (strcmp(type, "yarn") != 0 || factor <= 1.0) return; + + const double orig = js_num(d, js_get(d, rs, "original_max_position_embeddings"), 4096.0); + const double bf = js_num(d, js_get(d, rs, "beta_fast"), 32.0); + const double bs = js_num(d, js_get(d, rs, "beta_slow"), 1.0); + double low = floor(dim * log(orig / (bf * 2.0 * PI)) / (2.0 * log(base))); + double high = ceil(dim * log(orig / (bs * 2.0 * PI)) / (2.0 * log(base))); + if (low < 0.0) low = 0.0; + if (high > dim - 1) high = dim - 1; + if (low == high) high += 0.001; /* upstream's singularity guard */ + for (int j = 0; j < half; j++) { + double ramp = ((double)j - low) / (high - low); + ramp = ramp < 0.0 ? 0.0 : ramp > 1.0 ? 1.0 : ramp; + const double mask = 1.0 - ramp; /* 1 = extrapolate, 0 = interpolate */ + const double extra = c->rope_inv_freq[j]; + c->rope_inv_freq[j] = (float)((extra / factor) * (1.0 - mask) + extra * mask); + } + const double m_all = js_num(d, js_get(d, rs, "mscale_all_dim"), 0.0); + if (m_all != 0.0) { + const double ms = 0.1 * m_all * log(factor) + 1.0; + c->att_mul = (float)(ms * ms); + } +} + static void cfg_from_json(waste_config *c, const js_doc *d, int cfg) { c->n_layers = (int)js_int(d, js_get(d, cfg, "num_hidden_layers"), 0); @@ -892,6 +944,8 @@ static void cfg_from_json(waste_config *c, const js_doc *d, int cfg) js_str(d, js_at(d, a, 0), c->arch, sizeof c->arch); } + rope_init(c, d, cfg); + int lac = js_get(d, cfg, "linear_attn_config"); c->full_rank_gate = js_get(d, lac, "use_full_rank_gate") >= 0; c->gate_lower_bound = (float)js_num(d, js_get(d, lac, "gate_lower_bound"), 0.0); @@ -1037,6 +1091,15 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap, js_free(&d); free(src); return -2; /* -> WASTE_E_FORMAT */ } + /* rope_init leaves the table empty when the slice is wider than it can + * hold. Running anyway would apply no rotation, which is not a degraded + * result but an unordered one, so refuse instead. */ + if (!m->cfg.mla_nope && m->cfg.qk_rope > 2 * WASTE_MAX_ROPE_HALF) { + fprintf(stderr, "waste: qk_rope_head_dim %d needs rotation, this build " + "holds %d\n", m->cfg.qk_rope, 2 * WASTE_MAX_ROPE_HALF); + js_free(&d); free(src); + return -2; /* -> WASTE_E_FORMAT */ + } const waste_config *c = &m->cfg; int eq = js_get(&d, 0, "expert_quant"); @@ -2373,6 +2436,34 @@ static void kda_layer(waste_model *m, int L, const float *in, float *out) * the absorption, the scores, the softmax and the output projection — * the old expanded path ran that loop on one core. */ +/* Rotate one qk_rope-wide slice in place at `pos`. + * + * GPT-J / interleaved: pair j is (x[2j], x[2j+1]). Upstream reaches the same + * arithmetic by de-interleaving before a half-split rotate, + * q = q.view(b, h, s, d/2, 2).transpose(4, 3).reshape(b, h, s, d) + * so pairing dim j with dim j + qk_rope/2 instead — the LLaMA layout — rotates + * the wrong partners and still yields finite, weight-shaped output. + * + * The angles depend only on (pos, j), not on the head, so the caller builds + * the tables once per token per layer and every head reuses them. */ +static void rope_tables(const waste_config *c, int pos, float *cs, float *sn) +{ + for (int j = 0; j < c->qk_rope / 2; j++) { + const float a = (float)pos * c->rope_inv_freq[j]; + cs[j] = cosf(a); + sn[j] = sinf(a); + } +} + +static void rope_apply(int half, float *x, const float *cs, const float *sn) +{ + for (int j = 0; j < half; j++) { + const float e = x[2 * j], o = x[2 * j + 1]; + x[2 * j] = e * cs[j] - o * sn[j]; + x[2 * j + 1] = e * sn[j] + o * cs[j]; + } +} + typedef struct { waste_model *m; const waste_tensor *kvb; @@ -2454,8 +2545,21 @@ static void mla_layer(waste_model *m, int L, const float *in, float *out, int po in, c->kv_lora + c->qk_rope, hid); waste_rmsnorm(ckv, ckv, T(m, "%smodel.layers.%d.self_attn.kv_a_layernorm.weight", c->prefix, L), c->kv_lora, c->eps); - /* Cache the latent as-is — normalized kpass followed by the raw rope - * dims. kv_b_proj is not applied here at all; it is absorbed below. */ + /* Rotate before caching, not after: the cached entry is reused by every + * later query and carries this token's position, while the query carries + * the querying token's. Rotating on read would need the pair of positions + * and would redo the work once per (query, key). */ + if (!c->mla_nope) { + float cs[WASTE_MAX_ROPE_HALF], sn[WASTE_MAX_ROPE_HALF]; + rope_tables(c, pos, cs, sn); + const int half = c->qk_rope / 2; + for (int h = 0; h < nh; h++) + rope_apply(half, q + (size_t)h * qd + c->qk_nope, cs, sn); + rope_apply(half, ckv + c->kv_lora, cs, sn); + } + /* Cache the latent — normalized kpass followed by the rope dims, rotated + * unless the model is NoPE. kv_b_proj is not applied here at all; it is + * absorbed below. */ memcpy(m->latcache[L] + (size_t)pos * latd, ckv, (size_t)latd * sizeof(float)); /* WASTE_DUMP_LATENT=path appends the cached latent and the absorbed * query, the two things a KV-cache quantizer has to keep faithful. */ @@ -2475,7 +2579,9 @@ static void mla_layer(waste_model *m, int L, const float *in, float *out, int po a.S = m->n_kv[L]; a.qd = qd; a.qk_nope = c->qk_nope; a.qk_rope = c->qk_rope; a.vh = vh; a.kv_lora = c->kv_lora; a.latd = latd; - a.scale = 1.0f / sqrtf((float)qd); + /* YaRN raises the attention scale by mscale_all_dim^2 when the config + * sets it; att_mul is 1 otherwise, including on every NoPE model. */ + a.scale = c->att_mul / sqrtf((float)qd); waste_parallel_for(nh, 1, mla_head_range, &a); } if (c->mla_output_gate) { diff --git a/src/model.h b/src/model.h index 13352a4..c1f30c8 100644 --- a/src/model.h +++ b/src/model.h @@ -58,6 +58,11 @@ typedef struct { * across a hidden state, and one global scale would flatten the small * positions to zero. */ #define WASTE_VQ_LUT_BLK 32 + +/* Rotary pairs held per layer: qk_rope_head_dim / 2. 64 covers a 128-wide + * rope slice; every model in the family uses 64. A container needing + * rotation on a wider slice is refused at load rather than run unrotated. */ +#define WASTE_MAX_ROPE_HALF 64 int kda_layer[WASTE_MAX_LAYERS]; /* 1 if layer is KDA */ float eps, routed_scale; int renorm; @@ -81,6 +86,17 @@ typedef struct { * themselves model_type "kimi_linear", so this is the only field that * tells them apart by name rather than by feature. */ char arch[64]; + + /* --- rotary -------------------------------------------------------- */ + /* The Kimi models set mla_use_nope and are the reason this was absent: + * with NoPE the qk_rope dims pass through unrotated. Every DeepSeek-V3 + * model (V3, R1, K2) sets no such flag and needs the rotation, and in + * MLA those dims are the only positional signal — the nope dims are + * position-free by construction, so skipping it leaves attention unable + * to order the sequence. */ + int mla_nope; /* mla_use_nope: 1 = no rotation */ + float rope_inv_freq[WASTE_MAX_ROPE_HALF]; /* qk_rope/2 used, YaRN-adjusted */ + float att_mul; /* YaRN mscale^2 on the attn scale, 1 = none */ } waste_config; typedef struct { diff --git a/tools/deepseek_ref.py b/tools/deepseek_ref.py new file mode 100644 index 0000000..72134ac --- /dev/null +++ b/tools/deepseek_ref.py @@ -0,0 +1,331 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 SQLite Cloud, Inc. +""" +deepseek_ref.py — pure-PyTorch DeepSeek-V3 / Kimi-K2, running off a WASTE container. + +Companion to kimi_ref.py, which is the oracle for the Kimi-Linear family. That +one cannot serve K2: it hardcodes `linear_attn_config` (KeyError on a config +without KDA) and, more importantly, applies NO rotary at all -- + + # NoPE: mla_use_nope, so no rotary is applied to the "rot" dims + +which is correct for Kimi-Linear and K3 (`mla_use_nope: true`) and wrong for +every DeepSeek-V3 model. K2 sets no such flag and ships +`rope_theta: 50000` with YaRN scaling, so its `qk_rope_head_dim` dims must be +rotated. In MLA those dims are the ONLY positional signal -- the nope dims are +position-free by construction -- so omitting the rotation leaves the model +unable to order its own prompt. + +Same contract as kimi_ref.py: weights come FROM THE CONTAINER, trunk +dequantized on demand and experts dequantized per use, so a diff against the C +engine measures ARITHMETIC and not quantization error. Both sides see the same +3-bit experts. + + uv run --with torch python tools/deepseek_ref.py \ + --container /data/hermes/waste_containers/kimi-k2.waste \ + --ids 163594,14062,163601,... --top 10 + +`--no-rope --no-mscale` reproduces the engine exactly: layer 0's residual +stream matches its WASTE_DUMP_HIDDEN output to 0.000% rel L2 on K2. That +agreement is the reference's validation, so run it before reading any delta. + +Speed: a 61-layer forward dequantizes every routed expert it touches in +Python, and the distinct experts per layer grow with the token count — minutes +for a 15-token prompt, hours for a long one. Use the shortest prompt that +reproduces the behaviour under test. +""" + +import argparse +import json +import math +import os +import struct +import sys +import time + +import torch +import torch.nn.functional as F + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from kimi_ref import Container, rms_norm # noqa: E402 + + +# ------------------------------------------------------------------ yarn --- + +def yarn_find_correction_dim(num_rot, dim, base, max_pos): + return (dim * math.log(max_pos / (num_rot * 2 * math.pi))) / (2 * math.log(base)) + + +def yarn_get_mscale(scale, mscale): + return 1.0 if scale <= 1 else 0.1 * mscale * math.log(scale) + 1.0 + + +def rope_tables(cfg, dim): + """(inv_freq[dim/2], softmax_scale_multiplier), following + DeepseekV3YarnRotaryEmbedding in the checkpoint's modeling_deepseek.py. + + K2 carries beta_fast = beta_slow = 1.0 rather than HF's 32/1 defaults, + which collapses the correction range to dims 19..20: below that the + extrapolated frequency is kept, above it the frequency is interpolated by + 1/factor. YaRN rescales inv_freq globally, so it applies at every position, + including 0.""" + base = float(cfg.get("rope_theta", 10000.0)) + sc = cfg.get("rope_scaling") + half = torch.arange(0, dim, 2, dtype=torch.float32) / dim + freq_extra = 1.0 / (base ** half) + if not sc or sc.get("type") not in ("yarn",): + return freq_extra, 1.0 + factor = float(sc["factor"]) + orig = float(sc.get("original_max_position_embeddings", 4096)) + bf, bs = float(sc.get("beta_fast", 32)), float(sc.get("beta_slow", 1)) + freq_inter = freq_extra / factor + low = max(math.floor(yarn_find_correction_dim(bf, dim, base, orig)), 0) + high = min(math.ceil(yarn_find_correction_dim(bs, dim, base, orig)), dim - 1) + if low == high: + high += 0.001 # upstream's singularity guard + ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / (high - low)).clamp(0, 1) + mask = 1.0 - ramp # 1 => extrapolate, 0 => interpolate + inv_freq = freq_inter * (1 - mask) + freq_extra * mask + # cos/sin carry mscale / mscale_all_dim, which is 1.0 when the two are equal + # (K2: both 1.0). The attention scale carries mscale_all_dim squared, which + # is 1.8133x on K2. Same name, two different factors. + m_all = sc.get("mscale_all_dim", 0) + att_mul = yarn_get_mscale(factor, float(m_all)) ** 2 if m_all else 1.0 + return inv_freq, att_mul + + +def apply_rope(x, pos, inv_freq): + """Rotate the last dim of x [T, ..., dim] at integer positions `pos` [T]. + + GPT-J / interleaved convention: pair j is (x[2j], x[2j+1]). Upstream + de-interleaves before its half-split rotate, + + q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + and the two compose to exactly this. The LLaMA half-split form applied + directly to these weights pairs the wrong dims and still yields finite, + weight-shaped output.""" + ang = pos.float().unsqueeze(-1) * inv_freq # [T, dim/2] + cos, sin = ang.cos(), ang.sin() + shape = [x.shape[0]] + [1] * (x.dim() - 2) + [inv_freq.numel()] + cos, sin = cos.view(shape), sin.view(shape) + even, odd = x[..., 0::2], x[..., 1::2] + out = torch.empty_like(x) + out[..., 0::2] = even * cos - odd * sin + out[..., 1::2] = even * sin + odd * cos + return out + + +# ------------------------------------------------- row-subset dequant ------ + +def deq_rows(c, name, rows): + """Dequantize only `rows` of a trunk tensor. + + embed_tokens and lm_head are 163840 x 7168 on K2; materializing either in + f32 is 4.7 GB and the Q4G/Q8G unpack needs several times that transiently. + The forward needs a handful of embedding rows and can chunk the head, so + neither is ever built whole.""" + e = c._meta[name] + blob, shape = c._blob, e["shape"] + N = shape[-1] + if e["fmt"] == 0: + out = torch.empty(len(rows), N) + for i, r in enumerate(rows): + o = e["off"] + r * N * 4 + out[i] = torch.frombuffer(bytearray(blob[o:o + N * 4]), dtype=torch.float32) + return out + g = e["group"] + ng = (N + g - 1) // g + q4 = e["fmt"] == 3 + rb = ng * g // 2 if q4 else ng * g + out = torch.empty(len(rows), N) + for i, r in enumerate(rows): + o = e["off"] + r * rb + raw = bytearray(blob[o:o + rb]) + if q4: + b = torch.frombuffer(raw, dtype=torch.uint8).int() + v = (torch.stack([b & 0x0F, b >> 4], -1).view(ng, g) - 8).float() + else: + v = torch.frombuffer(raw, dtype=torch.int8).view(ng, g).float() + so = e["scale_off"] + r * ng * 2 + sc = torch.frombuffer(bytearray(blob[so:so + ng * 2]), + dtype=torch.float16).float().view(ng, 1) + out[i] = (v * sc).view(-1)[:N] + return out + + +# ----------------------------------------------------------------- model --- + +class DeepseekRef: + def __init__(self, c: Container, rope=True, verbose=False): + self.c, self.t, self.cfg = c, c.t, c.cfg + self.p = c.prefix + self.eps = self.cfg["rms_norm_eps"] + self.n_layers = self.cfg["num_hidden_layers"] + self.first_dense = self.cfg.get("first_k_dense_replace", 0) + self.use_rope = rope + self.verbose = verbose + self.qk_n = self.cfg["qk_nope_head_dim"] + self.qk_r = self.cfg["qk_rope_head_dim"] + self.inv_freq, self.att_mul = rope_tables(self.cfg, self.qk_r) + # A container this size makes the trunk cache the memory ceiling: one + # K2 layer is ~600 MB of f32 weights, so the 64 kimi_ref defaults to + # would hold tens of gigabytes. + c.t.cap = 6 + # No expert cache: one dequantized expert is three f32 matrices, 176 MB, + # and experts are never reused across layers because each layer has its + # own bank. A cache keyed on (layer, expert) would only grow — 512 + # entries want 90 GB. The grouping in moe() removes the redundant work + # instead, decoding each distinct expert once per layer. + + def mla(self, L, x, pos): + p = f"{self.p}model.layers.{L}.self_attn." + cfg, T = self.cfg, x.shape[0] + nh, qk_n, qk_r = cfg["num_attention_heads"], self.qk_n, self.qk_r + vh, qd = cfg["v_head_dim"], self.qk_n + self.qk_r + if cfg.get("q_lora_rank"): + qa = rms_norm(x @ self.t[p + "q_a_proj.weight"].T, + self.t[p + "q_a_layernorm.weight"], self.eps) + q = (qa @ self.t[p + "q_b_proj.weight"].T).view(T, nh, qd) + else: + q = (x @ self.t[p + "q_proj.weight"].T).view(T, nh, qd) + ckv = x @ self.t[p + "kv_a_proj_with_mqa.weight"].T + kpass, krot = ckv.split([cfg["kv_lora_rank"], qk_r], dim=-1) + kpass = rms_norm(kpass, self.t[p + "kv_a_layernorm.weight"], self.eps) + kb = (kpass @ self.t[p + "kv_b_proj.weight"].T).view(T, nh, qk_n + vh) + knope, val = kb.split([qk_n, vh], dim=-1) + + if self.use_rope: + qn, qr = q.split([qk_n, qk_r], dim=-1) + qr = apply_rope(qr, pos, self.inv_freq) + q = torch.cat([qn, qr], -1) + krot = apply_rope(krot, pos, self.inv_freq) + k = torch.cat([knope, krot.view(T, 1, qk_r).expand(T, nh, qk_r)], -1) + + scale = (qd ** -0.5) * self.att_mul + att = torch.einsum("thd,shd->hts", q, k) * scale + att = (att + torch.full((T, T), float("-inf")).triu(1)).softmax(-1) + o = torch.einsum("hts,shd->thd", att, val).reshape(T, nh * vh) + return o @ self.t[p + "o_proj.weight"].T + + def moe(self, L, x): + p = f"{self.p}model.layers.{L}.block_sparse_moe." + cfg, T = self.cfg, x.shape[0] + scores = torch.sigmoid(x.float() @ self.t[p + "gate.weight"].float().T) + choice = scores + self.t[p + "gate.e_score_correction_bias"].unsqueeze(0) + k = cfg["num_experts_per_token"] + idx = torch.topk(choice, k=k, dim=-1, sorted=False)[1] + w = scores.gather(1, idx) + if cfg.get("moe_renormalize", True): + w = w / (w.sum(-1, keepdim=True) + 1e-20) + w = w * cfg["routed_scaling_factor"] + + # Group tokens by expert so each distinct expert is decoded ONCE per + # layer rather than once per (token, slot). On a 15-token prompt that + # is ~90 decodes instead of 120, and the gap widens with length. + jobs = {} + for t in range(T): + for j in range(k): + jobs.setdefault(int(idx[t, j]), []).append((t, w[t, j])) + y = torch.zeros_like(x) + for eid, hits in jobs.items(): + E = self.c.expert(L, eid) + ts = [t for t, _ in hits] + xi = x[ts] + h = F.silu(xi @ E["gate"].T) * (xi @ E["up"].T) + o = h @ E["down"].T + for r, (t, wt) in enumerate(hits): + y[t] += wt * o[r] + sg, su, sd = (self.t[p + f"shared_experts.{n}.weight"] + for n in ("gate_proj", "up_proj", "down_proj")) + sh = F.silu(x @ sg.T) * (x @ su.T) + return y + sh @ sd.T + + def dense_mlp(self, L, x): + p = f"{self.p}model.layers.{L}.mlp." + h = F.silu(x @ self.t[p + "gate_proj.weight"].T) * (x @ self.t[p + "up_proj.weight"].T) + return h @ self.t[p + "down_proj.weight"].T + + def forward(self, ids, dump=None, upto=None): + pos = torch.arange(len(ids)) + x = deq_rows(self.c, self.p + "model.embed_tokens.weight", ids) + n = self.n_layers if upto is None else min(upto, self.n_layers) + for L in range(n): + pre = f"{self.p}model.layers.{L}." + t0 = time.time() + x = x + self.mla(L, rms_norm(x, self.t[pre + "input_layernorm.weight"], self.eps), pos) + h = rms_norm(x, self.t[pre + "post_attention_layernorm.weight"], self.eps) + x = x + (self.dense_mlp(L, h) if L < self.first_dense else self.moe(L, h)) + if dump: + with open(dump, "ab" if L else "wb") as f: + v = x[-1].float().tolist() + f.write(struct.pack(f"<{len(v)}f", *v)) + if self.verbose: + print(f" layer {L:>3}/{n} {time.time()-t0:6.1f}s " + f"|x|={x[-1].norm():.3f}", + flush=True) + if upto is not None: + return None + x = rms_norm(x, self.t[self.p + "model.norm.weight"], self.eps)[-1] + # lm_head in row blocks: 163840 x 7168 is 4.7 GB in f32 and we only + # need the resulting vector of logits. + name = self.p + "lm_head.weight" + V = self.c._meta[name]["shape"][0] + out = torch.empty(V) + B = 8192 + for beg in range(0, V, B): + rows = list(range(beg, min(beg + B, V))) + out[beg:beg + len(rows)] = deq_rows(self.c, name, rows) @ x + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--container", required=True) + ap.add_argument("--ids", help="comma-separated token ids (from `waste tokenize`)") + ap.add_argument("--top", type=int, default=10) + ap.add_argument("--no-rope", action="store_true", + help="skip the rotary, i.e. what the C engine does today") + ap.add_argument("--no-mscale", action="store_true", + help="drop YaRN's mscale^2 from the attention scale, which " + "the engine also omits; with --no-rope this reproduces " + "the engine") + ap.add_argument("--dump-hidden", help="per-layer residual stream, engine's format") + ap.add_argument("--upto", type=int, help="stop after N layers (bisecting)") + ap.add_argument("--threads", type=int, + help="torch intra-op threads; default is one per physical core") + ap.add_argument("-v", "--verbose", action="store_true") + a = ap.parse_args() + + if a.threads: + torch.set_num_threads(a.threads) + ids = [int(x) for x in a.ids.replace(" ", ",").split(",") if x] + c = Container(a.container) + m = DeepseekRef(c, rope=not a.no_rope, verbose=a.verbose) + if a.no_mscale: + m.att_mul = 1.0 + print(f"container {a.container}", file=sys.stderr) + print(f"rope {'OFF (engine behaviour)' if a.no_rope else 'ON'}" + f" att_mul {m.att_mul:.4f} layers {m.n_layers} ntok {len(ids)}", + file=sys.stderr) + t0 = time.time() + lg = m.forward(ids, dump=a.dump_hidden, upto=a.upto) + if lg is None: + print(f"stopped after {a.upto} layers, hidden dumped", file=sys.stderr) + return 0 + pr = lg.softmax(-1) + top = torch.topk(lg, a.top) + print(json.dumps({ + "prompt_tokens": len(ids), + "rope": not a.no_rope, + "elapsed_s": round(time.time() - t0, 1), + "top": [{"id": int(i), "logit": round(float(v), 4), + "prob": round(float(pr[i]), 6)} + for v, i in zip(top.values, top.indices)], + })) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 90b111d39b20b30bcd277abfc284863212ed1aa7 Mon Sep 17 00:00:00 2001 From: Marco Bambini Date: Fri, 7 Aug 2026 07:54:58 +0200 Subject: [PATCH 2/4] Cover the rotary: a container shaped like a V3, and three checks over it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The suite stayed green through the whole window in which `src/` applied no rotary, and it would have stayed green after a fix that pairs the wrong dims. Both have the same cause: every container the suite can reach is a Kimi, every Kimi sets `mla_use_nope`, and so nothing in `tests/` ever entered `rope_init` or `rope_apply`. This is the missing half of the previous commit. `make_test_container.py --rope` writes a DeepSeek-V3 at the 1/18 scale the file already builds a Kimi-Linear at: no `mla_use_nope`, `rope_theta` and the YaRN block copied from Kimi-K2-Instruct's config, and no `linear_attn_config` at all, which is what makes every layer MLA. All-MLA is deliberate twice over — it exercises the rotation at depth rather than in the single full-attention layer the Kimi mix leaves, and it is the shape `deepseek_ref.py` can read, since not indexing `linear_attn_config` is exactly what separates it from `kimi_ref.py`. K2's rope block rather than V3's because `beta_fast == beta_slow == 1.0` collapses YaRN's correction range to a two-dim ramp, which is the more awkward of the two to get right. The checks build their own container instead of using `$MODEL`, so they run on every host and do not wait on weights nobody can convert yet — #26 is what makes a real V3 container, and the shape is what the engine branches on. - rotated MLA against the PyTorch oracle - chunked prefill == token-at-a-time with rotation, which holds by construction today because `mla_layer` is per-token on both paths, and is exactly the "by construction" a later batched MLA would break quietly - a rope slice wider than `WASTE_MAX_ROPE_HALF` is refused at load The first takes the same two-source shape as the Kimi oracle above it: generate from `deepseek_ref.py` where `uv` exists, fall back to a fixture where it does not, so the Linux image without `uv` runs it rather than skipping it. Unlike that one the fixture ships, because this container is generated rather than converted and so is byte-reproducible at `--seed 0` — the sidecar carries a digest of the container it was made from, so a later change to the generator's weights reads as "regenerate me" and not as an engine bug. The fixture is the reference's logits, never the engine's. `deepseek_ref.py` grows the `--dump` that `kimi_ref.py` already had, so the diff is over whole logit vectors and not a printed top-k. Verified by reverting `src/model.c` and `src/model.h` to their pre-fix state with `tests/` and `tools/` left alone: the oracle check and the refusal check both fail, which is the property that makes them worth having. Both fallback paths were exercised directly — `uv` off `PATH` passes against the fixture, and a corrupted digest skips with the regenerate message instead of reporting a divergence. `set -o pipefail` sank the refusal check on the first run, because a refused load exits non-zero and that is the point; the output is read into a variable now, with a comment saying why. Suite on this commit: 46 passed, 0 failed, 2 skipped against Kimi-Linear and K3 (43/0/2 before), 39/0/9 on the synthetic path CI takes, and `make asan` 33/0/14. Fuzzer and the 168 serve checks unchanged. Co-Authored-By: Claude Opus 5 --- tests/fixtures/oracle_ropesynth_16tok.bin | Bin 0 -> 1024 bytes tests/fixtures/oracle_ropesynth_16tok.json | 8 ++ tests/run.sh | 113 +++++++++++++++++++++ tools/deepseek_ref.py | 8 ++ tools/make_test_container.py | 39 ++++++- 5 files changed, 167 insertions(+), 1 deletion(-) create mode 100644 tests/fixtures/oracle_ropesynth_16tok.bin create mode 100644 tests/fixtures/oracle_ropesynth_16tok.json diff --git a/tests/fixtures/oracle_ropesynth_16tok.bin b/tests/fixtures/oracle_ropesynth_16tok.bin new file mode 100644 index 0000000000000000000000000000000000000000..6abda04948c73064b6d76208d8b45b1a84addfe2 GIT binary patch literal 1024 zcmV~$4LHyV|z1s_>hz!Q}J){o3K@I`~y&{tK- z^v_pg+q{z)S}w=qj%J!Ryc3nvf5W^oGgbNslsW%0rkHF{;%5~;@O*&E)E(lhj1@dE z*a95|E$qxSjbL3AO+hzKi)BYcNViYUQg;`Ln$%40nQF(E44(t9aT$4qh4P(gC+J}C z9{lra9(-Xjfll+xKV&vR@?;Zp?ufvKrFQ6iy8s4kk#FCnK>c_x^e6sTujsgp;Zuj8 zjtiu4`kL+C7tAefD$&8xPx<b^SFvZ6{QsOt0i*GLchFeb9 z@n8RHC%?B9VAQ-jpN^3vVT=$1ONim=F zizc6K?ldYmQLS_)j4kMg=Fh&Qo|FGjkM0$B*X{x_H6B_WE72;fN8b-6bY83FwYTEw zPRlBsH_fkM>nR&@8TO>2Zhz`bZvnRkcV^Q2f_1I~+Rl~HRhbH1-^xZs_W`>6s1=f< zze0O$HuG|R$PFJ?!l?EbB%1Y*b1RCCcgj&Wa4KLRppRYDHqmSGCW^|ih~^SG)#Ms^ zS7|2Fq2+K_GbPTA_2HWES%K%=WU%$Udhw{*%(v7&h216o_+1)`cCyphXZB#utdRDDVN*IY(en$S*ARwgJ?O%Wz9y5avtgl!fJ7u%$|DM zPEz~BTjD0YeS>7)I&2SD@FkOraj1f!y%fQ>H;JlVw2J4}&0sAlbHG8R7N3mXgT~$z ziho`W&3Q$zws}}=G{o|D>ls+FYfKom*41wl7<>Qs9fY;jLRjPw=*y&_RJ&4Q`F1?? ZdLbU$WE4h27qA!)C-g3?ptR*%`2X(d_`?7I literal 0 HcmV?d00001 diff --git a/tests/fixtures/oracle_ropesynth_16tok.json b/tests/fixtures/oracle_ropesynth_16tok.json new file mode 100644 index 0000000..9382359 --- /dev/null +++ b/tests/fixtures/oracle_ropesynth_16tok.json @@ -0,0 +1,8 @@ +{ + "container": "python3 tools/make_test_container.py --rope --seed 0 ", + "container_sha256": "7ae7f6c06a6c8956b668d334c2b14ae9a7d94e9adc1bbbc89db3c30563b6510b", + "ids": "3,7,11,5,9,13,2,17,4,8,19,23,6,29,12,31", + "oracle": "uv run --with torch --no-project python tools/deepseek_ref.py --container --ids --dump tests/fixtures/oracle_ropesynth_16tok.bin", + "what": "Last token's logits, f32, vocab 256 \u2014 the layout test_forward writes. Computed by the PyTorch reference, not by the engine, so the diff means something.", + "why_the_digest": "Unlike the Kimi-Linear fixture this container is generated, not converted, so it is byte-reproducible at seed 0 and the fixture can ship. The digest is over that container: change make_test_container.py's weights and the fixture is stale, which has to read as 'regenerate me' and not as an engine bug." +} \ No newline at end of file diff --git a/tests/run.sh b/tests/run.sh index cf35765..430c78b 100755 --- a/tests/run.sh +++ b/tests/run.sh @@ -631,6 +631,119 @@ else sk "engine checks" "no container at $MODEL" fi +# --------------------------------------------------------------- rotary ---- +head_ "rotary (MLA on a model that is not NoPE)" + +# Everything above this point runs on a Kimi, and every Kimi sets +# mla_use_nope — so none of it reaches rope_init or rope_apply in +# src/model.c. The rotation was absent from the engine for that reason and +# the suite stayed green throughout, which is the failure this section +# exists to stop repeating. +# +# It builds its own DeepSeek-V3-shaped container rather than using $MODEL, +# so it runs on every host and does not depend on which weights happen to be +# on disk. Nobody ships a V3 container yet — that needs the fp8 reader — +# but the shape is what the engine branches on, and the shape is free. +ROPE="$TMP/rope.waste" +RIDS=3,7,11,5,9,13,2,17,4,8,19,23,6,29,12,31 +if ! python3 tools/make_test_container.py --rope --seed 0 "$ROPE" >/dev/null 2>&1; then + sk "rotary checks" "make_test_container.py --rope did not build a container" +else + ./test_forward "$ROPE" "$RIDS" "$TMP/rope_seq.bin" 0 >/dev/null 2>&1 + if [ ! -s "$TMP/rope_seq.bin" ]; then + no "the engine did not run a container without mla_use_nope" + else + # Same two-source shape as the Kimi oracle above: generate from the + # reference where torch is available, fall back to the fixture where + # it is not. This container is *generated* rather than converted, so + # unlike that one it is byte-reproducible at seed 0 and the fixture + # is portable — the digest below is what says so. + RGEN="" + if command -v uv >/dev/null 2>&1; then + uv run --no-project --with torch \ + python tools/deepseek_ref.py --container "$ROPE" --ids "$RIDS" \ + --dump "$TMP/rope_ref.bin" >/dev/null 2>&1 || true + [ -s "$TMP/rope_ref.bin" ] && RGEN="$TMP/rope_ref.bin" + fi + RFIX=tests/fixtures/oracle_ropesynth_16tok.bin + rope_why="" + if [ -z "$RGEN" ] && [ -f "${RFIX%.bin}.json" ]; then + rope_why=$(python3 - "$ROPE" "${RFIX%.bin}.json" <<'PY' +import hashlib, json, os, sys +h = hashlib.sha256() +for n in sorted(os.listdir(sys.argv[1])): + h.update(n.encode()) + h.update(open(os.path.join(sys.argv[1], n), "rb").read()) +want = json.load(open(sys.argv[2])).get("container_sha256") +if want and h.hexdigest() != want: + print("no uv to generate one, and make_test_container.py --rope no " + "longer builds the container this fixture was made from — " + "regenerate it, see " + os.path.basename(sys.argv[2])) +PY +) + fi + if [ -n "$rope_why" ]; then + sk "engine vs the rotary oracle" "$rope_why" + elif [ -n "$RGEN" ] || [ -f "$RFIX" ]; then + if python3 - "$TMP/rope_seq.bin" "${RGEN:-$RFIX}" <<'PY' +import struct, sys +def L(p): + b = open(p, "rb").read() + return struct.unpack(f"<{len(b)//4}f", b) +a, b = L(sys.argv[1]), L(sys.argv[2]) +sys.exit(0 if max(abs(x - y) for x, y in zip(a, b)) < 1e-3 else 1) +PY + then + if [ -n "$RGEN" ] + then ok "rotated MLA matches a PyTorch oracle built from this container" + else ok "rotated MLA matches the shipped rotary fixture" + fi + # An engine that skips the rotation still produces finite, + # weight-shaped logits — that is why this went unnoticed — so the + # diff is the only thing that separates the two. + else no "rotated MLA diverges from the oracle" + fi + else + sk "engine vs the rotary oracle" \ + "no fixture; regenerate with tools/deepseek_ref.py --dump" + fi + + # The chunked check above runs on $MODEL, which is NoPE. mla_layer is + # per-token on both paths, so this should hold by construction — and + # it is exactly the kind of "by construction" that a later batched + # MLA would break silently. + WASTE_CHUNK=1 ./test_forward "$ROPE" "$RIDS" "$TMP/rope_chunk.bin" 0 >/dev/null 2>&1 + if python3 - "$TMP/rope_seq.bin" "$TMP/rope_chunk.bin" <<'PY' +import struct, sys +def L(p): + b = open(p, "rb").read() + return struct.unpack(f"<{len(b)//4}f", b) +a, b = L(sys.argv[1]), L(sys.argv[2]) +d = max(abs(x - y) for x, y in zip(a, b)) +sys.exit(0 if d < 1e-3 and a.index(max(a)) == b.index(max(b)) else 1) +PY + then ok "chunked prefill == token-at-a-time with rotation" + else no "chunked prefill diverges on a rotated model" + fi + fi + + # The rope table is a fixed WASTE_MAX_ROPE_HALF pairs. A container that + # needs more must be refused at load: running it would apply no rotation, + # and that is not a degraded answer but an unordered one. + WIDE="$TMP/rope_wide.waste" + if ! python3 tools/make_test_container.py --rope --qk-rope 132 "$WIDE" >/dev/null 2>&1; then + sk "a rope slice wider than the build holds is refused" "container not built" + # Read into a variable rather than piping: a refused load is a non-zero + # exit, which is the point, and under `set -o pipefail` that would sink + # the pipeline no matter what grep found. + elif printf '%s' "$(./test_forward "$WIDE" 3,7,11 "$TMP/wide.bin" 0 2>&1 || true)" \ + | grep -q "needs rotation"; then + ok "a rope slice wider than the build holds is refused at load" + else + no "an over-wide rope slice loaded instead of being refused" + fi +fi + # --------------------------------------------------------------- budget ---- head_ "RAM budget" diff --git a/tools/deepseek_ref.py b/tools/deepseek_ref.py index 72134ac..f6e7d0a 100644 --- a/tools/deepseek_ref.py +++ b/tools/deepseek_ref.py @@ -292,6 +292,9 @@ def main(): "the engine also omits; with --no-rope this reproduces " "the engine") ap.add_argument("--dump-hidden", help="per-layer residual stream, engine's format") + ap.add_argument("--dump", default="", + help="last token's logits as f32, the layout test_forward " + "writes — this is what tests/run.sh diffs against") ap.add_argument("--upto", type=int, help="stop after N layers (bisecting)") ap.add_argument("--threads", type=int, help="torch intra-op threads; default is one per physical core") @@ -314,6 +317,11 @@ def main(): if lg is None: print(f"stopped after {a.upto} layers, hidden dumped", file=sys.stderr) return 0 + if a.dump: + v = lg.float().tolist() + with open(a.dump, "wb") as f: + f.write(struct.pack(f"<{len(v)}f", *v)) + print(f"dumped logits -> {a.dump}", file=sys.stderr) pr = lg.softmax(-1) top = torch.topk(lg, a.top) print(json.dumps({ diff --git a/tools/make_test_container.py b/tools/make_test_container.py index b1f4838..0216b7b 100644 --- a/tools/make_test_container.py +++ b/tools/make_test_container.py @@ -82,6 +82,23 @@ } C_KDA = H_KDA * D_KDA +# --rope turns the above into a DeepSeek-V3 at the same scale, which is the +# only shape that reaches src/model.c's rotary: the Kimi models set +# mla_use_nope and pass the qk_rope dims through unrotated, so a container +# built from CFG as it stands leaves rope_init and rope_apply dead. +# +# The rope block is Kimi-K2-Instruct's config.json verbatim. DeepSeek-V3 and +# R1 ship the same shape with factor 40 and beta_fast 32; K2's beta_fast == +# beta_slow == 1.0 is the more awkward of the two because it collapses YaRN's +# correction range to a two-dim ramp, so it is the one worth pinning. +V3_ROPE = { + "rope_theta": 50000.0, + "rope_scaling": {"beta_fast": 1.0, "beta_slow": 1.0, "factor": 32.0, + "mscale": 1.0, "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn"}, +} + def f32(vals): return struct.pack("<%df" % len(vals), *vals) @@ -262,11 +279,31 @@ def main(): help="put the text tensors under a tensor_prefix, e.g. " "language_model., and add one tensor outside it — " "K3's shape, and the one the loader skips") + ap.add_argument("--rope", action="store_true", + help="a DeepSeek-V3 instead of a Kimi-Linear: every layer " + "MLA, no mla_use_nope, and rope_theta with YaRN — the " + "only shape that reaches the engine's rotary") + ap.add_argument("--qk-rope", type=int, metavar="N", + help="override qk_rope_head_dim. With --rope, a slice " + "wider than the build's WASTE_MAX_ROPE_HALF pair " + "table has to be refused at load, not run unrotated") args = ap.parse_args() rng = random.Random(args.seed) os.makedirs(args.out, exist_ok=True) cfg = dict(CFG) + if args.rope: + # Dropping linear_attn_config is what makes every layer MLA, so the + # rotation is exercised at depth rather than in the one full-attention + # layer the Kimi mix leaves. It also makes the container readable by + # tools/deepseek_ref.py, which — unlike kimi_ref.py — does not index + # linear_attn_config and does apply the rotary. + del cfg["mla_use_nope"], cfg["linear_attn_config"] + cfg["model_type"] = "deepseek_v3" + cfg["architectures"] = ["DeepseekV3ForCausalLM"] + cfg.update(V3_ROPE) + if args.qk_rope: + cfg["qk_rope_head_dim"] = args.qk_rope if args.tokenizer: # Every special has to be a real row of the embedding table and the # head: a container whose vocab_size stops short of its own specials @@ -280,7 +317,7 @@ def main(): qd = cfg["qk_nope_head_dim"] + cfg["qk_rope_head_dim"] kvl, rope = cfg["kv_lora_rank"], cfg["qk_rope_head_dim"] moe, dense = cfg["moe_intermediate_size"], cfg["intermediate_size"] - kda = {l - 1 for l in cfg["linear_attn_config"]["kda_layers"]} + kda = {l - 1 for l in cfg.get("linear_attn_config", {}).get("kda_layers", [])} t = Trunk(rng, args.prefix) t.quant("model.embed_tokens.weight", [cfg["vocab_size"], hid]) From b8c872c9077c573c1497439a9df7bdf0b7dc57e7 Mon Sep 17 00:00:00 2001 From: fab2s Date: Fri, 7 Aug 2026 21:54:50 +0200 Subject: [PATCH 3/4] Read mla_use_nope by value, and refuse the rope shapes this does not implement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two ways the rotation could still be skipped with nobody told. `mla_use_nope` was read by presence, which is the idiom cfg_from_json uses for its other flags. A container carrying it as `false` therefore loaded as NoPE and produced exactly the pre-fix result: unrotated logits, finite and weight-shaped, no error. convert.py copies the source config verbatim, so a checkpoint that writes the flag out rather than omitting it reaches that path. js_bool() reads the token instead — `false` is false, a missing key is still the default. The other flags keep the presence idiom, where a misread costs a feature rather than the sequence order. A rope_scaling shape the ramp does not implement fell through to plain RoPE just as quietly. Any `type` but yarn landed there, and so did mscale != mscale_all_dim, where upstream puts the ratio on cos/sin and this does not. Nothing shipping is affected — V3, R1 and K2 all set `type: yarn` with both mscales at 1.0 — but the argument for refusing an over-wide slice is the same argument, so rope_init leaves a reason in cfg.rope_err and the load refuses on it. The over-wide check moves into that mechanism, leaving one refusal point rather than a condition re-derived at the call site. Two shapes are read rather than refused. `rope_type` is HF's rename of `type`, so it is an alias: a config that only spells it that way is rotated, not turned away. And factor <= 1 stays a fall-through, because YaRN's ramp is the identity there and both mscales collapse to 1 — plain RoPE is the right answer, not a degraded one. tools/deepseek_ref.py refuses the same two shapes rather than approximating them, so it remains an oracle for exactly what the engine accepts. Three checks over make_test_container.py's new knobs: a container built at the same seed with `mla_use_nope: false` gives logits byte-identical to the same model without the key, and each unimplemented rope_scaling is refused at load. 37 passed, 0 failed, 12 skipped on the synthetic path (34/0/12 before), `make asan` 36/0/13 with the six rotary checks passing and no sanitizer report, fuzzer 400 cases 0 crashed. A full Kimi-K2 container still loads and reports unchanged. --- src/json.h | 8 +++++ src/model.c | 63 +++++++++++++++++++++++++++--------- src/model.h | 3 ++ tests/run.sh | 60 +++++++++++++++++++++++++--------- tools/deepseek_ref.py | 16 +++++++-- tools/make_test_container.py | 18 +++++++++++ 6 files changed, 135 insertions(+), 33 deletions(-) diff --git a/src/json.h b/src/json.h index 2b863dc..3efe95d 100644 --- a/src/json.h +++ b/src/json.h @@ -208,6 +208,14 @@ static inline int64_t js_int(const js_doc *d, int t, int64_t dflt) return (int64_t)v; } +/* Reads a value, not a presence: `false` is false, and anything that is not + * a JSON boolean — including a missing key — is dflt. */ +static inline int js_bool(const js_doc *d, int t, int dflt) +{ + if (t < 0 || t >= d->n || d->tok[t].type != JS_BOOL) return dflt; + return d->src[d->tok[t].start] == 't'; +} + /* Copies at most cap-1 bytes; always NUL-terminates. */ static inline const char *js_str(const js_doc *d, int t, char *buf, size_t cap) { diff --git a/src/model.c b/src/model.c index 7d8928f..fa6a5de 100644 --- a/src/model.c +++ b/src/model.c @@ -858,25 +858,60 @@ static int cfg_sane(const waste_config *c) * mscale_all_dim SQUARED, which is 1.8133x on K2. * - YaRN rescales inv_freq globally, so it applies from position 0. It is * not a long-context-only correction that a short prompt can ignore. + * + * A shape this does not implement leaves a reason in c->rope_err and the + * load refuses on it. Falling through to plain RoPE instead would be the + * same failure this function was added to fix: not a degraded answer but an + * unordered one, and one that looks like weight-shaped logits. */ static void rope_init(waste_config *c, const js_doc *d, int cfg) { const double PI = 3.14159265358979323846; c->att_mul = 1.0f; - c->mla_nope = js_get(d, cfg, "mla_use_nope") >= 0; + c->rope_err[0] = 0; + /* By value, not by presence: a container carrying "mla_use_nope": false + * has to rotate. The presence idiom used for the other flags costs a + * feature when it misreads; here it costs the sequence order. */ + c->mla_nope = js_bool(d, js_get(d, cfg, "mla_use_nope"), 0); const int dim = c->qk_rope, half = dim / 2; - if (c->mla_nope || half <= 0 || half > WASTE_MAX_ROPE_HALF) return; + if (c->mla_nope || half <= 0) return; + if (half > WASTE_MAX_ROPE_HALF) { + snprintf(c->rope_err, sizeof c->rope_err, + "qk_rope_head_dim %d needs rotation, this build holds %d", + dim, 2 * WASTE_MAX_ROPE_HALF); + return; + } const double base = js_num(d, js_get(d, cfg, "rope_theta"), 10000.0); for (int j = 0; j < half; j++) c->rope_inv_freq[j] = (float)(1.0 / pow(base, (double)(2 * j) / dim)); const int rs = js_get(d, cfg, "rope_scaling"); - if (rs < 0) return; - char type[16]; - js_str(d, js_get(d, rs, "type"), type, sizeof type); + if (rs < 0) return; /* plain RoPE, computed above */ + char type[24]; + int ty = js_get(d, rs, "type"); + if (ty < 0) ty = js_get(d, rs, "rope_type"); /* HF renamed the key */ + js_str(d, ty, type, sizeof type); + if (strcmp(type, "yarn") != 0) { + snprintf(c->rope_err, sizeof c->rope_err, + "rope_scaling type \"%s\" is not implemented, only yarn", type); + return; + } + /* factor <= 1 is not a refusal: YaRN's ramp is the identity there and + * both mscales collapse to 1, so plain RoPE is the right answer. */ const double factor = js_num(d, js_get(d, rs, "factor"), 1.0); - if (strcmp(type, "yarn") != 0 || factor <= 1.0) return; + if (factor <= 1.0) return; + /* Unequal mscales put a ratio on cos/sin that nothing here applies. + * V3, R1, K2 and V2 all ship them equal; HF's defaults (1 and 0) are + * not, so an omitted mscale_all_dim lands here too. */ + const double m_one = js_num(d, js_get(d, rs, "mscale"), 1.0); + const double m_dim = js_num(d, js_get(d, rs, "mscale_all_dim"), 0.0); + if (m_one != m_dim) { + snprintf(c->rope_err, sizeof c->rope_err, + "rope_scaling mscale %g != mscale_all_dim %g, and the ratio " + "on cos/sin is not implemented", m_one, m_dim); + return; + } const double orig = js_num(d, js_get(d, rs, "original_max_position_embeddings"), 4096.0); const double bf = js_num(d, js_get(d, rs, "beta_fast"), 32.0); @@ -893,9 +928,8 @@ static void rope_init(waste_config *c, const js_doc *d, int cfg) const double extra = c->rope_inv_freq[j]; c->rope_inv_freq[j] = (float)((extra / factor) * (1.0 - mask) + extra * mask); } - const double m_all = js_num(d, js_get(d, rs, "mscale_all_dim"), 0.0); - if (m_all != 0.0) { - const double ms = 0.1 * m_all * log(factor) + 1.0; + if (m_dim != 0.0) { + const double ms = 0.1 * m_dim * log(factor) + 1.0; c->att_mul = (float)(ms * ms); } } @@ -1091,12 +1125,11 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap, js_free(&d); free(src); return -2; /* -> WASTE_E_FORMAT */ } - /* rope_init leaves the table empty when the slice is wider than it can - * hold. Running anyway would apply no rotation, which is not a degraded - * result but an unordered one, so refuse instead. */ - if (!m->cfg.mla_nope && m->cfg.qk_rope > 2 * WASTE_MAX_ROPE_HALF) { - fprintf(stderr, "waste: qk_rope_head_dim %d needs rotation, this build " - "holds %d\n", m->cfg.qk_rope, 2 * WASTE_MAX_ROPE_HALF); + /* rope_init leaves no table for a shape it does not implement. Running + * anyway would apply no rotation, which is not a degraded result but an + * unordered one, so refuse instead. */ + if (m->cfg.rope_err[0]) { + fprintf(stderr, "waste: %s\n", m->cfg.rope_err); js_free(&d); free(src); return -2; /* -> WASTE_E_FORMAT */ } diff --git a/src/model.h b/src/model.h index c1f30c8..03547ba 100644 --- a/src/model.h +++ b/src/model.h @@ -97,6 +97,9 @@ typedef struct { int mla_nope; /* mla_use_nope: 1 = no rotation */ float rope_inv_freq[WASTE_MAX_ROPE_HALF]; /* qk_rope/2 used, YaRN-adjusted */ float att_mul; /* YaRN mscale^2 on the attn scale, 1 = none */ + char rope_err[128]; /* non-empty: a shape rope_init does not + * implement, and why. The load refuses on + * it rather than running unrotated. */ } waste_config; typedef struct { diff --git a/tests/run.sh b/tests/run.sh index 430c78b..2481fb6 100755 --- a/tests/run.sh +++ b/tests/run.sh @@ -725,23 +725,53 @@ PY then ok "chunked prefill == token-at-a-time with rotation" else no "chunked prefill diverges on a rotated model" fi - fi - # The rope table is a fixed WASTE_MAX_ROPE_HALF pairs. A container that - # needs more must be refused at load: running it would apply no rotation, - # and that is not a degraded answer but an unordered one. - WIDE="$TMP/rope_wide.waste" - if ! python3 tools/make_test_container.py --rope --qk-rope 132 "$WIDE" >/dev/null 2>&1; then - sk "a rope slice wider than the build holds is refused" "container not built" - # Read into a variable rather than piping: a refused load is a non-zero - # exit, which is the point, and under `set -o pipefail` that would sink - # the pipeline no matter what grep found. - elif printf '%s' "$(./test_forward "$WIDE" 3,7,11 "$TMP/wide.bin" 0 2>&1 || true)" \ - | grep -q "needs rotation"; then - ok "a rope slice wider than the build holds is refused at load" - else - no "an over-wide rope slice loaded instead of being refused" + # Same model, same seed, one line of config: mla_use_nope written out + # as false instead of omitted. A loader that tests the key for + # presence reads that as NoPE and skips the rotation, which is the + # pre-fix engine — so these logits have to match the ones above. + FALSE="$TMP/rope_nopefalse.waste" + if ! python3 tools/make_test_container.py --rope --nope-false --seed 0 \ + "$FALSE" >/dev/null 2>&1; then + sk "mla_use_nope: false rotates" "container not built" + else + ./test_forward "$FALSE" "$RIDS" "$TMP/rope_false.bin" 0 >/dev/null 2>&1 + if [ -s "$TMP/rope_false.bin" ] && cmp -s "$TMP/rope_seq.bin" "$TMP/rope_false.bin" + then ok "mla_use_nope: false rotates, like the same model without the key" + else no "mla_use_nope: false was read as NoPE and skipped the rotation" + fi + fi fi + + # Shapes rope_init does not implement. Each has to be refused at load: + # running one would apply no rotation or the wrong one, and that is not a + # degraded answer but an unordered one. + rope_refused() { # + local what=$1 want=$2; shift 2 + local dir="$TMP/rope_bad.waste" + rm -rf "$dir" + if ! python3 tools/make_test_container.py --rope "$@" "$dir" >/dev/null 2>&1; then + sk "$what is refused at load" "container not built" + # Read into a variable rather than piping: a refused load is a + # non-zero exit, which is the point, and under `set -o pipefail` that + # would sink the pipeline no matter what grep found. + elif printf '%s' "$(./test_forward "$dir" 3,7,11 "$TMP/bad.bin" 0 2>&1 || true)" \ + | grep -q "$want"; then + ok "$what is refused at load" + else + no "$what loaded instead of being refused" + fi + } + + # The rope table is a fixed WASTE_MAX_ROPE_HALF pairs. + rope_refused "a rope slice wider than the build holds" \ + "needs rotation" --qk-rope 132 + # Anything but yarn — linear, dynamic — reaches none of the ramp below it. + rope_refused "an unimplemented rope_scaling type" \ + "not implemented, only yarn" --rope-type linear + # Unequal mscales put a ratio on cos/sin that rope_tables does not apply. + rope_refused "rope_scaling with mscale != mscale_all_dim" \ + "not implemented" --mscale 0.707 fi # --------------------------------------------------------------- budget ---- diff --git a/tools/deepseek_ref.py b/tools/deepseek_ref.py index f6e7d0a..3db7c82 100644 --- a/tools/deepseek_ref.py +++ b/tools/deepseek_ref.py @@ -74,8 +74,12 @@ def rope_tables(cfg, dim): sc = cfg.get("rope_scaling") half = torch.arange(0, dim, 2, dtype=torch.float32) / dim freq_extra = 1.0 / (base ** half) - if not sc or sc.get("type") not in ("yarn",): + kind = sc.get("type", sc.get("rope_type")) if sc else None + if not sc: return freq_extra, 1.0 + if kind != "yarn": + raise SystemExit(f"rope_scaling type {kind!r} is not implemented here; " + "the engine refuses the same shape at load") factor = float(sc["factor"]) orig = float(sc.get("original_max_position_embeddings", 4096)) bf, bs = float(sc.get("beta_fast", 32)), float(sc.get("beta_slow", 1)) @@ -89,8 +93,14 @@ def rope_tables(cfg, dim): inv_freq = freq_inter * (1 - mask) + freq_extra * mask # cos/sin carry mscale / mscale_all_dim, which is 1.0 when the two are equal # (K2: both 1.0). The attention scale carries mscale_all_dim squared, which - # is 1.8133x on K2. Same name, two different factors. - m_all = sc.get("mscale_all_dim", 0) + # is 1.8133x on K2. Same name, two different factors. Unequal mscales are + # refused rather than approximated, so this stays an oracle for exactly the + # shapes the engine accepts. + m_one, m_all = sc.get("mscale", 1.0), sc.get("mscale_all_dim", 0) + if float(m_one) != float(m_all): + raise SystemExit(f"rope_scaling mscale {m_one} != mscale_all_dim " + f"{m_all}; the ratio on cos/sin is not implemented " + "here, and the engine refuses it at load") att_mul = yarn_get_mscale(factor, float(m_all)) ** 2 if m_all else 1.0 return inv_freq, att_mul diff --git a/tools/make_test_container.py b/tools/make_test_container.py index 0216b7b..d08009d 100644 --- a/tools/make_test_container.py +++ b/tools/make_test_container.py @@ -287,6 +287,18 @@ def main(): help="override qk_rope_head_dim. With --rope, a slice " "wider than the build's WASTE_MAX_ROPE_HALF pair " "table has to be refused at load, not run unrotated") + ap.add_argument("--nope-false", action="store_true", + help="write mla_use_nope: false rather than omitting it. " + "Same model either way — a loader that tests the key " + "for presence reads it as NoPE and skips the rotation") + ap.add_argument("--rope-type", metavar="T", + help="override rope_scaling.type, e.g. linear — a scaling " + "the engine does not implement has to be refused, not " + "quietly run as plain RoPE") + ap.add_argument("--mscale", type=float, metavar="X", + help="override rope_scaling.mscale, leaving mscale_all_dim " + "at 1.0. Unequal mscales put a ratio on cos/sin that " + "the engine does not apply, so it refuses instead") args = ap.parse_args() rng = random.Random(args.seed) os.makedirs(args.out, exist_ok=True) @@ -302,6 +314,12 @@ def main(): cfg["model_type"] = "deepseek_v3" cfg["architectures"] = ["DeepseekV3ForCausalLM"] cfg.update(V3_ROPE) + if args.nope_false: + cfg["mla_use_nope"] = False + if args.rope_type: + cfg["rope_scaling"] = dict(cfg["rope_scaling"], type=args.rope_type) + if args.mscale is not None: + cfg["rope_scaling"] = dict(cfg["rope_scaling"], mscale=args.mscale) if args.qk_rope: cfg["qk_rope_head_dim"] = args.qk_rope if args.tokenizer: From 7604709a07f607ebc301f0672a85337ddcac2790 Mon Sep 17 00:00:00 2001 From: fab2s Date: Sat, 8 Aug 2026 23:10:01 +0200 Subject: [PATCH 4/4] Load a rope_scaling of null, and refuse an mla_use_nope that is not a boolean MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `"rope_scaling": null` was refused. js_get returns a token for a JSON null but it is not a JS_OBJ, so `type` and `rope_type` both came back absent, js_str left the buffer empty and the type refusal fired — on a container whose correct plain-RoPE inv_freq the loop above had already built and then threw away. null is how an HF config says "no scaling" and convert.py copies configs verbatim, so it is the shape most containers on disk carry; the Kimi ones are saved from it only by returning early on mla_use_nope. js_size is 0 for a null and for {} alike, so one condition covers both. tools/deepseek_ref.py already read this shape correctly, so the two agree again. An mla_use_nope that is present but not a boolean took js_bool's default and rotated. Nothing writes `1` or `"true"`, but defaulting picks the sequence order out of a manifest that never said which one it wanted, and picks it silently — the failure this file exists to remove. js_typeof() tells "absent" from "present but not that", and present-but-not-boolean now leaves a reason in rope_err the way an unimplemented rope_scaling does. A rope_scaling carrying no type says so, rather than reporting the type as the empty string. That is the message someone debugs from. --nope takes a JSON value instead of being a --nope-false switch, and --rope-scaling builds the null / {} / absent / no-type shapes. Three checks over them: null and {} load with logits byte-identical to the same container with no key at all, and the other two are refused at load. 40 passed, 0 failed, 12 skipped on the synthetic path (37/0/12 before), `make asan` 39/0/13 with all nine rotary checks passing and no sanitizer report, fuzzer 400 cases 0 crashed 0 hung. The shipped rotary fixture's container digest is unchanged, so CI's no-uv path still compares. A full Kimi-K2 container still opens and reports unchanged. --- src/json.h | 8 ++++++++ src/model.c | 29 +++++++++++++++++++++++++---- tests/run.sh | 28 +++++++++++++++++++++++++++- tools/make_test_container.py | 28 ++++++++++++++++++++++------ 4 files changed, 82 insertions(+), 11 deletions(-) diff --git a/src/json.h b/src/json.h index 3efe95d..0fe6225 100644 --- a/src/json.h +++ b/src/json.h @@ -188,6 +188,14 @@ static inline int js_at(const js_doc *d, int arr, int i) return p; } +/* The token's type, or -1 for no such token. The accessors below fold + * "absent" and "present but not that type" into the same default; a caller + * that has to tell those apart needs this. */ +static inline int js_typeof(const js_doc *d, int t) +{ + return (t < 0 || t >= d->n) ? -1 : (int)d->tok[t].type; +} + static inline double js_num(const js_doc *d, int t, double dflt) { if (t < 0 || t >= d->n || d->tok[t].type != JS_NUM) return dflt; diff --git a/src/model.c b/src/model.c index fa6a5de..231247c 100644 --- a/src/model.c +++ b/src/model.c @@ -871,8 +871,19 @@ static void rope_init(waste_config *c, const js_doc *d, int cfg) c->rope_err[0] = 0; /* By value, not by presence: a container carrying "mla_use_nope": false * has to rotate. The presence idiom used for the other flags costs a - * feature when it misreads; here it costs the sequence order. */ - c->mla_nope = js_bool(d, js_get(d, cfg, "mla_use_nope"), 0); + * feature when it misreads; here it costs the sequence order. + * + * Present but not a boolean is refused rather than defaulted, because + * defaulting picks the sequence order from a manifest that did not say + * which one it wanted, and picks it silently. */ + const int nope = js_get(d, cfg, "mla_use_nope"); + c->mla_nope = 0; + if (nope >= 0 && js_typeof(d, nope) != JS_BOOL) { + snprintf(c->rope_err, sizeof c->rope_err, + "mla_use_nope is present but is not true or false"); + return; + } + c->mla_nope = js_bool(d, nope, 0); const int dim = c->qk_rope, half = dim / 2; if (c->mla_nope || half <= 0) return; if (half > WASTE_MAX_ROPE_HALF) { @@ -886,12 +897,22 @@ static void rope_init(waste_config *c, const js_doc *d, int cfg) for (int j = 0; j < half; j++) c->rope_inv_freq[j] = (float)(1.0 / pow(base, (double)(2 * j) / dim)); + /* A key that is absent, null or {} all mean no scaling, and js_size is 0 + * for each — the plain-RoPE table above is already the whole answer. + * null is how HF configs spell it and convert.py copies them verbatim, + * so this is the common shape, not the corner. */ const int rs = js_get(d, cfg, "rope_scaling"); - if (rs < 0) return; /* plain RoPE, computed above */ + if (rs < 0 || js_size(d, rs) == 0) return; char type[24]; int ty = js_get(d, rs, "type"); if (ty < 0) ty = js_get(d, rs, "rope_type"); /* HF renamed the key */ - js_str(d, ty, type, sizeof type); + js_str(d, ty, type, sizeof type); /* "" if absent or not a string */ + if (!type[0]) { + snprintf(c->rope_err, sizeof c->rope_err, + "rope_scaling carries no type string, and only yarn is " + "implemented"); + return; + } if (strcmp(type, "yarn") != 0) { snprintf(c->rope_err, sizeof c->rope_err, "rope_scaling type \"%s\" is not implemented, only yarn", type); diff --git a/tests/run.sh b/tests/run.sh index 2481fb6..decc721 100755 --- a/tests/run.sh +++ b/tests/run.sh @@ -731,7 +731,7 @@ PY # presence reads that as NoPE and skips the rotation, which is the # pre-fix engine — so these logits have to match the ones above. FALSE="$TMP/rope_nopefalse.waste" - if ! python3 tools/make_test_container.py --rope --nope-false --seed 0 \ + if ! python3 tools/make_test_container.py --rope --nope false --seed 0 \ "$FALSE" >/dev/null 2>&1; then sk "mla_use_nope: false rotates" "container not built" else @@ -741,6 +741,25 @@ PY else no "mla_use_nope: false was read as NoPE and skipped the rotation" fi fi + + # null is how an HF config says "no scaling" and convert.py copies it + # verbatim, so it is the shape most containers on disk carry. It has + # to load as plain RoPE — the same as {} and the same as no key at + # all — rather than being read as a scaling with an unknown type. + none_ok=1 + for shape in drop null empty; do + dir="$TMP/rope_$shape.waste" + rm -rf "$dir" + python3 tools/make_test_container.py --rope --rope-scaling "$shape" \ + --seed 0 "$dir" >/dev/null 2>&1 || { none_ok=0; break; } + ./test_forward "$dir" "$RIDS" "$TMP/rs_$shape.bin" 0 >/dev/null 2>&1 + [ -s "$TMP/rs_$shape.bin" ] || { none_ok=0; break; } + cmp -s "$TMP/rs_drop.bin" "$TMP/rs_$shape.bin" || { none_ok=0; break; } + done + if [ "$none_ok" = 1 ] + then ok "rope_scaling null and {} load as plain RoPE, like no key at all" + else no "rope_scaling null or {} did not load as plain RoPE" + fi fi # Shapes rope_init does not implement. Each has to be refused at load: @@ -772,6 +791,13 @@ PY # Unequal mscales put a ratio on cos/sin that rope_tables does not apply. rope_refused "rope_scaling with mscale != mscale_all_dim" \ "not implemented" --mscale 0.707 + # A scaling object that carries no type is not the same as no scaling. + rope_refused "rope_scaling that carries no type" \ + "carries no type" --rope-scaling notype + # Present but not a boolean names no sequence order, so neither does a + # default picked for it. + rope_refused "mla_use_nope that is not true or false" \ + "not true or false" --nope 1 fi # --------------------------------------------------------------- budget ---- diff --git a/tools/make_test_container.py b/tools/make_test_container.py index d08009d..3572aa6 100644 --- a/tools/make_test_container.py +++ b/tools/make_test_container.py @@ -287,10 +287,18 @@ def main(): help="override qk_rope_head_dim. With --rope, a slice " "wider than the build's WASTE_MAX_ROPE_HALF pair " "table has to be refused at load, not run unrotated") - ap.add_argument("--nope-false", action="store_true", - help="write mla_use_nope: false rather than omitting it. " - "Same model either way — a loader that tests the key " - "for presence reads it as NoPE and skips the rotation") + ap.add_argument("--nope", metavar="JSON", + help="write mla_use_nope with this JSON value rather than " + "omitting the key. `false` is the same model, and a " + "loader that tests for presence reads it as NoPE and " + "skips the rotation; `1` or `\"true\"` say nothing a " + "loader may act on, and have to be refused") + ap.add_argument("--rope-scaling", choices=("null", "empty", "drop", "notype"), + metavar="SHAPE", + help="replace the YaRN block: null | empty ({}) | drop " + "(no key) all mean no scaling and must load as plain " + "RoPE; notype is an object with a factor and no type, " + "which must be refused") ap.add_argument("--rope-type", metavar="T", help="override rope_scaling.type, e.g. linear — a scaling " "the engine does not implement has to be refused, not " @@ -314,12 +322,20 @@ def main(): cfg["model_type"] = "deepseek_v3" cfg["architectures"] = ["DeepseekV3ForCausalLM"] cfg.update(V3_ROPE) - if args.nope_false: - cfg["mla_use_nope"] = False + if args.nope is not None: + cfg["mla_use_nope"] = json.loads(args.nope) if args.rope_type: cfg["rope_scaling"] = dict(cfg["rope_scaling"], type=args.rope_type) if args.mscale is not None: cfg["rope_scaling"] = dict(cfg["rope_scaling"], mscale=args.mscale) + if args.rope_scaling == "null": + cfg["rope_scaling"] = None + elif args.rope_scaling == "empty": + cfg["rope_scaling"] = {} + elif args.rope_scaling == "drop": + del cfg["rope_scaling"] + elif args.rope_scaling == "notype": + cfg["rope_scaling"] = {"factor": 40.0, "beta_fast": 1.0, "beta_slow": 1.0} if args.qk_rope: cfg["qk_rope_head_dim"] = args.qk_rope if args.tokenizer: