diff --git a/src/json.h b/src/json.h index 2b863dc..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; @@ -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) { diff --git a/src/model.c b/src/model.c index 51ef35b..231247c 100644 --- a/src/model.c +++ b/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) { diff --git a/src/model.h b/src/model.h index 13352a4..03547ba 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,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 { diff --git a/tests/fixtures/oracle_ropesynth_16tok.bin b/tests/fixtures/oracle_ropesynth_16tok.bin new file mode 100644 index 0000000..6abda04 Binary files /dev/null and b/tests/fixtures/oracle_ropesynth_16tok.bin differ 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..decc721 100755 --- a/tests/run.sh +++ b/tests/run.sh @@ -631,6 +631,175 @@ 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 + + # 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 + + # 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: + # 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 + # 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 ---- head_ "RAM budget" diff --git a/tools/deepseek_ref.py b/tools/deepseek_ref.py new file mode 100644 index 0000000..3db7c82 --- /dev/null +++ b/tools/deepseek_ref.py @@ -0,0 +1,349 @@ +#!/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) + 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)) + 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. 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 + + +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("--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") + 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 + 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({ + "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()) diff --git a/tools/make_test_container.py b/tools/make_test_container.py index b1f4838..3572aa6 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,65 @@ 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") + 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 " + "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) 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.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: # 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 +351,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])