Apply rotary to MLA's rope dims on models that are not NoPE

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;
@@ -208,6 +216,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)
 {

src/model.c

@@ -848,6 +848,113 @@ 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.
+ *
+ * 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->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.
+     *
+     * 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) {
+        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));
+
+    /* 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 || 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);              /* "" 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);
+        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 (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);
+    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);
+    }
+    if (m_dim != 0.0) {
+        const double ms = 0.1 * m_dim * 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 +999,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 +1146,14 @@ 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 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 */
+    }
     const waste_config *c = &m->cfg;
 
     int eq = js_get(&d, 0, "expert_quant");
@@ -2373,6 +2490,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 +2599,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 +2633,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) {

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,20 @@ 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 */
+    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 {

tests/fixtures/oracle_ropesynth_16tok.json

@@ -0,0 +1,8 @@
+{
+ "container": "python3 tools/make_test_container.py --rope --seed 0 <dir>",
+ "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 <dir> --ids <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."
+}