From 3843b0d705eceb7772a6f0c48f585edeb5571498 Mon Sep 17 00:00:00 2001 From: Luca Barbato Date: Sun, 11 Oct 2026 00:23:37 +0200 Subject: [PATCH 1/2] =?UTF-8?q?kolibri1=20TT:=20localize=20the=20L1-moe=20?= =?UTF-8?q?3%=20divergence=20to=20router-gain=20amplification=20of=20the?= =?UTF-8?q?=20upstream=20bf16=20drift=20class=20=E2=80=94=20and=20fix=20th?= =?UTF-8?q?e=20f32-out=20GEMM=20contract=20the=20instrument=20caught?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The layer-1 moe divergence after the slot-dequant fix is NOT a stride/staging/accounting defect. The teacher-forced instrument (stage bit-hashes, the router dump on both arms, the router weight hash probe, host-reference logits for both arms' router GEMMs, the element-level L0 dhn diff) proves the chain: the first diverging quantity is the L0 input rms_norm output at step 0 (1 bf16 ULP class); by L0 dh2 the drift carries outlier element diffs (~0.4% rms, 3 ULP on a 48-magnitude element); the router gate weight is byte-identical across arms and BOTH router GEMMs match host reference logits computed from their own inputs (CPU 1.4e-5, TT 0.15 max) — and the kolibri router row has ||w||_2 ~ 1 over K=2560 against logits of scale O(5), so that sub-ULP-class drift injects O(1-5) ABSOLUTE logit noise and flips the near-tie top-6-of-384 selection at every layer including L0. The ~3% L1 moe divergence IS the flipped mixture. Classification: selection flips DRIVEN by continuous precision drift — the upstream adjudication decision, not a local patch. The instrument did catch one contract defect, fixed here: the TT kMatmul/kMatmulBT default produced a BF16 tensor upcast on commit for f32-out callers — all 384 router logits sat on the bf16 grid (bf16grid=384/384) with bf16 partial sums, violating the CPU row's LinearBTRaw f32 contract. f32-out now requests dtype FLOAT32 plus fp32_dest_acc_en; post-fix bf16grid=0/384. The device token gate (ON arm, clean reset) keeps 4 HARD flips at the same positions (gaps -5.62 / -3.85 / -14.76 / -3.75 vs -5.92 / -4.56 / -13.35 / -3.06) and the same deterministic memoized-expert DRAM exhaustion abort: the contract fix does not clear the flips because the drift class drives them. Recorded in ISSUE-LOCAL-01M4ER0E9HHM95YZYJB7T5FECN with the full evidence chain; the OFF arm was not re-gated (it aborts on the same OOM before completing). FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai/glm-5.3-flash [maki] --- .../ISSUE-LOCAL-01M4ER0E9HHM95YZYJB7T5FECN.md | 43 ++++ .../models/kolibri1_forward.cpp | 133 +++++++++++- .../models/kolibri1_tt_forward.cpp | 202 +++++++++++++++++- src/vt/tenstorrent/tenstorrent_ops.cpp | 36 +++- 4 files changed, 393 insertions(+), 21 deletions(-) diff --git a/.agents/issues/MODEL-TEXT-kolibri-1-tenstorrent/ISSUE-LOCAL-01M4ER0E9HHM95YZYJB7T5FECN.md b/.agents/issues/MODEL-TEXT-kolibri-1-tenstorrent/ISSUE-LOCAL-01M4ER0E9HHM95YZYJB7T5FECN.md index 19e5d8f91f..b8adb9fe94 100644 --- a/.agents/issues/MODEL-TEXT-kolibri-1-tenstorrent/ISSUE-LOCAL-01M4ER0E9HHM95YZYJB7T5FECN.md +++ b/.agents/issues/MODEL-TEXT-kolibri-1-tenstorrent/ISSUE-LOCAL-01M4ER0E9HHM95YZYJB7T5FECN.md @@ -155,3 +155,46 @@ B2b-i landed the dense-resident device forward with the routed-expert tier delib slot-pool `what()=="1"` throw did NOT fire on the fixed substrate in any instrument or gate run (0 occurrences). The remaining flips after the fix are the residual bf16 device-arm drift class, not the byte bug. + +- 2026-10-10 (branch row/tt-kolibri-l1moe, off 519c466b8): the L1-moe ~3% + relative divergence is LOCALIZED and CLASSIFIED as the drift class, not a + stride/staging bug — instrument on the fixed substrate, teacher-forced + 'der Mond...' (VT_TF_DBG_PROMPT=1, 12 forwards), OFF-vs-CPU with per-stage + bf16 bit hashes (StageDump hash=/sq=), the router dump (RouteDump, both + arms), the router weight hash probe, host-reference logits for BOTH arms' + router GEMMs, and the element-level L0 dhn diff. Evidence chain: (1) the + FIRST diverging quantity is the L0 input rms_norm output at step 0 + (hash b3eeb298 vs 2b64e8d9; element diff 1 bf16 ULP class, top ~0.03-0.06 + on values to 14, rms 0.002) — the adjudicated substrate drift, not new; + (2) by L0 dh2 (the moe input) the drift carries OUTLIER element diffs + (max 47.25 vs 48.0, 3 bf16 ULP on a 48-magnitude element; sq 2.166e4 vs + 2.175e4, ~0.4% rms, diff-norm ~10 over 2560 elems) from the device + attention/norm chain; (3) the router gate weight is BYTE-IDENTICAL across + arms (hash e6512148 at L0) and BOTH arms' router GEMMs match host + reference logits computed from their own inputs (CPU 1.4e-5, TT 0.15 + max) — the GEMMs are correct; (4) the amplification is the ROUTER GAIN: + the kolibri router row has ||w||_2 ~ 1 over K=2560 against logits of + scale O(5), so a 0.4%-rms input drift injects O(1-5) ABSOLUTE logit noise + (measured max 4.8-7.0 per layer, rel_rms 0.3-1.4) and the top-6-of-384 + sigmoid-logit-add selection — a near-tie boundary — flips 2-4 of the 6 + experts AT EVERY LAYER INCLUDING L0. The ~3% L1 moe divergence IS the + flipped mixture. NO local defect: staging byte-exact (pool readback, + prior record), dequants byte-exact, weights identical, GEMMs verified + against host references. ONE CONTRACT DEFECT found and fixed by the + instrument (kept, red-first bf16grid evidence): the TT kMatmul/kMatmulBT + default produced a BF16 tensor upcast on commit for f32-out callers — + all 384 router logits sat on the bf16 grid (bf16grid=384/384) with bf16 + partial sums, violating the CPU row's LinearBTRaw f32 contract + (tenstorrent_ops.cpp: f32 out now requests dtype FLOAT32 + + fp32_dest_acc_en; post-fix bf16grid=0/384). GATE (ON arm, clean reset, + this branch, post-fix): 4 HARD flips at the SAME positions as pre-fix + ('der Mond...' step 7 gap -5.92 -> -5.62; 'Australia' step 5 -4.56 -> + -3.85; 'Wissen ist Macht' step 8 -13.35 -> -14.76; 'x1 = 3...' step 22 + -3.06 -> -3.75), then the same deterministic memoized-expert DRAM + exhaustion abort — the contract fix does NOT clear the flips because the + driver is the upstream drift class. CLASSIFICATION: selection flips + DRIVEN by continuous bf16 device-arm precision drift through the + attention/norm chain; per the row's adjudication rule this is an + upstream adjudication decision (f32-enabling the device attention/norm + chain or accepting the band), not a local patch. ISSUE STAYS OPEN on the + device gates and the OOM debt. diff --git a/src/vllm/model_executor/models/kolibri1_forward.cpp b/src/vllm/model_executor/models/kolibri1_forward.cpp index aace3c2daf..0ac82cc15a 100644 --- a/src/vllm/model_executor/models/kolibri1_forward.cpp +++ b/src/vllm/model_executor/models/kolibri1_forward.cpp @@ -36,6 +36,7 @@ #include #include #include +#include #include #include #include @@ -49,6 +50,11 @@ #include "vllm/model_executor/models/host_parallel.h" // the ONE pool (#1664) #include "vt/ops.h" +// SCRATCH DEBUG (layer-1 moe localization): the CPU arm's last router logits, +// read by the TT arm's RouteDump for the direct per-expert diff. +std::vector g_kolibri1_dbg_route_logits; +std::vector g_kolibri1_dbg_dhn; + namespace vllm { using dense_attn::DBuf; @@ -307,14 +313,22 @@ DBuf AttentionBlock(Dev d, const Kolibri1AttnWeights& w, const Kolibri1Params& p }; const int64_t n = t * hq * dh; for (int64_t i = 0; i < n; ++i) s += val(bf[i]); - std::fprintf(stderr, "[STAGE] pa_out sum=%.6f first=%.6f\n", s, val(bf[0])); + std::fprintf(stderr, "[STAGE][CPU] pa_out sum=%.6f first=%.6f\n", s, val(bf[0])); } return LinearBT(d, o_in, w.o_proj, t); // [T, H] } // ── MoE block (kolibri1.py:126-207) ───────────────────────────────────────── + +// SCRATCH DEBUG (layer-1 moe localization, defined after StageDump below; +// forward-declared here because MoeBlock calls it): +void StageDump(const char* tag, int64_t layer, DBuf& buf, Dev d); +void RouteDump(int64_t layer, const std::vector& logits, + const Kolibri1HostRouting& route, int64_t t, int64_t e, + int64_t top_k); + DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, - const Tensor& dhn, int64_t t) { + const Tensor& dhn, int64_t t, int64_t layer) { const int64_t h = p.hidden_size; const int64_t e = p.num_experts; const int64_t top_k = p.num_experts_per_tok; @@ -336,11 +350,68 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, be.Synchronize(d.q); } + if (std::getenv("VT_KOLIBRI1_TT_STAGE_DUMP") != nullptr) { + static std::set seen; + if (!seen.count(layer)) { + seen.insert(layer); + Tensor bt3 = ResidentWeight(d, w.router_gate, {e, h}); + const auto* bf3 = reinterpret_cast(bt3.data); + uint32_t hh = 2166136261u; + double ss = 0; + for (int64_t i = 0; i < e * h; ++i) { + hh = (hh ^ bf3[i]) * 16777619u; + uint32_t u = static_cast(bf3[i]) << 16; + float f; + std::memcpy(&f, &u, 4); + ss += f; + } + std::fprintf(stderr, "[STAGE][CPU] router_gate hash=%08x sum=%.4f\n", + hh, ss); + // Host reference logits from the same dhn + router_gate bytes. + const int64_t hk2 = w.router_gate.shape[1]; + std::vector xb(static_cast(t) * hk2 * 2); + std::memcpy(xb.data(), dhn.data, xb.size()); + const auto* xb16 = reinterpret_cast(xb.data()); + std::vector href(static_cast(t) * e, 0.f); + for (int64_t n0 = 0; n0 < e; ++n0) { + const auto* wr = + reinterpret_cast(bt3.data) + + static_cast(n0) * hk2; + for (int64_t i = 0; i < t; ++i) { + double acc = 0; + for (int64_t k = 0; k < hk2; ++k) { + auto dec = [](uint16_t b) { + uint32_t u = static_cast(b) << 16; + float f; + std::memcpy(&f, &u, 4); + return static_cast(f); + }; + acc += dec(wr[k]) * dec(xb16[static_cast(i) * hk2 + k]); + } + href[static_cast(i) * e + n0] = static_cast(acc); + } + } + double madd = 0, mlog = 0; + for (size_t k = 0; k < href.size(); ++k) { + madd = std::max(madd, std::fabs(static_cast(href[k]) - + static_cast(logits[k]))); + mlog = std::max( + mlog, + std::fabs(static_cast(href[k]) - + static_cast(g_kolibri1_dbg_route_logits.empty() + ? 0.f + : href[k]))); + } + std::fprintf(stderr, "[STAGE][CPU] router_href maxdiff_dev=%.6f\n", + madd); + } + } Kolibri1HostRouting route; { prof::Scope prof("moe_topk"); route = SigmoidLogitAddRouting(logits, bias, t, e, top_k); } + RouteDump(layer, logits, route, t, e, top_k); DBuf dtw(d, DType::kF32, {t, top_k}, const_cast(route.weights.data())); @@ -349,6 +420,7 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, // Shared expert: UNGATED, always added (:146-188). DBuf shared = ExpertMlp(d, w.shared_experts, dhn, t, p.shared_expert_intermediate_size); + StageDump("moe_shared", layer, shared, d); // Routed experts: the CPU-reference per-expert gather/scatter loop (the // mimo-v2 MoeBlock pattern). @@ -381,6 +453,13 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, } DBuf o = ExpertMlp(d, w.experts[static_cast(ex)], gathered.t(), ne, inter); + { + char tag[32]; + std::snprintf(tag, sizeof tag, "moe_gath#%d", static_cast(ex)); + StageDump(tag, layer, gathered, d); + std::snprintf(tag, sizeof tag, "moe_ex#%d", static_cast(ex)); + StageDump(tag, layer, o, d); + } for (int64_t i = 0; i < ne; ++i) { d.b.Copy(d.q, static_cast(expert_out.ptr()) + @@ -417,8 +496,9 @@ void StageDump(const char* tag, int64_t layer, DBuf& buf, Dev d) { std::vector tmp(static_cast(n) * 2); buf.Download(d, tmp.data()); const auto* bf = reinterpret_cast(tmp.data()); - double sum = 0; + double sum = 0, sq = 0; float mx = 0; + uint32_t hash = 2166136261u; auto val = [](uint16_t bits) { uint32_t u = static_cast(bits) << 16; float f; @@ -428,12 +508,45 @@ void StageDump(const char* tag, int64_t layer, DBuf& buf, Dev d) { for (int64_t i = 0; i < n; ++i) { const float v = val(bf[i]); sum += v; + sq += static_cast(v) * v; if (std::fabs(v) > mx) mx = std::fabs(v); + hash = (hash ^ bf[i]) * 16777619u; } - std::fprintf(stderr, "[STAGE] L%lld %s sum=%.6f max=%.6f first=[%.6f %.6f %.6f %.6f]\n", - static_cast(layer), tag, sum, mx, val(bf[0]), val(bf[1]), + std::fprintf(stderr, "[STAGE][CPU] L%lld %s sum=%.6f sq=%.6e hash=%08x max=%.6f first=[%.6f %.6f %.6f %.6f]\n", + static_cast(layer), tag, sum, sq, hash, mx, val(bf[0]), val(bf[1]), val(bf[2]), val(bf[3])); } +// SCRATCH DEBUG (layer-1 moe localization, same env as StageDump): the +// router's raw logits, top-k selection and weights per token — the +// discrete-selection leg of the OFF-vs-CPU stage diff. +void RouteDump(int64_t layer, const std::vector& logits, + const Kolibri1HostRouting& route, int64_t t, int64_t e, + int64_t top_k) { + if (std::getenv("VT_KOLIBRI1_TT_STAGE_DUMP") == nullptr) return; + for (int64_t i = 0; i < t; ++i) { + double ls = 0; + for (int64_t j = 0; j < e; ++j) + ls += logits[static_cast(i * e + j)]; + std::string ids, ws; + for (int64_t k = 0; k < top_k; ++k) { + const size_t idx = static_cast(i * top_k + k); + ids += (k ? "," : "") + std::to_string(route.ids[idx]); + ws += (k ? "," : "") + + [](float v) { + char b[32]; + std::snprintf(b, sizeof b, "%.6f", v); + return std::string(b); + }(route.weights[idx]); + } + std::fprintf(stderr, + "[STAGE][CPU] L%lld route tok=%lld logit_sum=%.6f ids=[%s] " + "w=[%s]\n", + static_cast(layer), static_cast(i), ls, + ids.c_str(), ws.c_str()); + } + g_kolibri1_dbg_route_logits = logits; +} + } // namespace @@ -515,7 +628,7 @@ ForwardLogits ForwardKolibri1Forward( } } for (int64_t i = 0; i < h; ++i) s += val(bf[static_cast(i)]); - std::fprintf(stderr, "[STAGE] gam sum=%.6f first=[%.6f %.6f %.6f %.6f]\n", + std::fprintf(stderr, "[STAGE][CPU] gam sum=%.6f first=[%.6f %.6f %.6f %.6f]\n", s, val(bf[0]), val(bf[1]), val(bf[2]), val(bf[3])); } Tensor dhn_t = dhn.t(); @@ -532,6 +645,12 @@ ForwardLogits ForwardKolibri1Forward( // Attention -> post_attn_norm (no residual). StageDump("dhn", l, dhn, d); + if (l == 0 && std::getenv("VT_KOLIBRI1_TT_STAGE_DUMP") != nullptr) { + const int64_t hn = dhn.t().shape[0] * dhn.t().shape[1]; + g_kolibri1_dbg_dhn.assign( + reinterpret_cast(dhn.t().data), + reinterpret_cast(dhn.t().data) + hn); + } DBuf attn = AttentionBlock(d, lw.attn, p, lw.is_sliding, dhn.t(), si.positions.t(), si, *kv_ptr, t); StageDump("attn", l, attn, d); @@ -561,7 +680,7 @@ ForwardLogits ForwardKolibri1Forward( StageDump("dh2", l, dh2, d); // MoE on EVERY layer -> post_ffn_norm (no residual). - DBuf moe = MoeBlock(d, lw.moe, p, dh2.t(), t); + DBuf moe = MoeBlock(d, lw.moe, p, dh2.t(), t, l); StageDump("moe", l, moe, d); DBuf moe_n(d, DType::kBF16, {t, h}); Tensor w_pf = ResidentWeight(d, lw.post_ffn_norm, {h}); diff --git a/src/vllm/model_executor/models/kolibri1_tt_forward.cpp b/src/vllm/model_executor/models/kolibri1_tt_forward.cpp index 9c05b1e271..fb6d06f631 100644 --- a/src/vllm/model_executor/models/kolibri1_tt_forward.cpp +++ b/src/vllm/model_executor/models/kolibri1_tt_forward.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -42,6 +43,11 @@ #include "vllm/model_executor/models/host_parallel.h" // the ONE pool (#1664) #include "vt/ops.h" +// SCRATCH DEBUG (layer-1 moe localization): the CPU arm's stashed router +// logits (defined in kolibri1_forward.cpp, same process), read by RouteDump. +extern std::vector g_kolibri1_dbg_route_logits; +extern std::vector g_kolibri1_dbg_dhn; + namespace vllm { using dense_attn::DBuf; @@ -89,6 +95,10 @@ bool StageDumpOn() { return on; } +void RouteDump(int64_t layer, const std::vector& logits, + const Kolibri1HostRouting& route, int64_t t, int64_t e, + int64_t top_k); + void StageDump(const char* tag, int64_t layer, DBuf& buf, Dev d) { if (!StageDumpOn()) return; const Tensor& t = buf.t(); @@ -96,8 +106,9 @@ void StageDump(const char* tag, int64_t layer, DBuf& buf, Dev d) { std::vector tmp(static_cast(n) * 2); buf.Download(d, tmp.data()); const auto* bf = reinterpret_cast(tmp.data()); - double sum = 0; + double sum = 0, sq = 0; float mx = 0; + uint32_t hash = 2166136261u; auto val = [](uint16_t bits) { uint32_t u = static_cast(bits) << 16; float f; @@ -107,13 +118,83 @@ void StageDump(const char* tag, int64_t layer, DBuf& buf, Dev d) { for (int64_t i = 0; i < n; ++i) { const float v = val(bf[i]); sum += v; + sq += static_cast(v) * v; if (std::fabs(v) > mx) mx = std::fabs(v); + hash = (hash ^ bf[i]) * 16777619u; } - std::fprintf(stderr, "[STAGE] L%lld %s sum=%.6f max=%.6f first=[%.6f %.6f %.6f %.6f]\n", - static_cast(layer), tag, sum, mx, val(bf[0]), val(bf[1]), + std::fprintf(stderr, "[STAGE][TT] L%lld %s sum=%.6f sq=%.6e hash=%08x max=%.6f first=[%.6f %.6f %.6f %.6f]\n", + static_cast(layer), tag, sum, sq, hash, mx, val(bf[0]), val(bf[1]), val(bf[2]), val(bf[3])); } +// SCRATCH DEBUG: dump a keyed weight once per layer (first forward only). +bool l3dbg(int64_t layer) { + static std::set seen; + if (seen.count(layer)) return false; + seen.insert(layer); + return true; +} + +// SCRATCH DEBUG (layer-1 moe localization): the router's raw logits, top-k +// selection and weights per token — mirrors the CPU arm's RouteDump (same +// env, same line shape) so the discrete-selection leg diffs arm vs arm. The +// CPU arm stashes its logits in g_kolibri1_dbg_route_logits (same process); +// this side prints the per-element diff against it plus the bf16-grid +// membership count (a grid hit means the logit was computed IN bf16). +float Kolibri1DbgToBf16(float v) { + uint32_t u; + std::memcpy(&u, &v, 4); + u = ((u + 0x7fffu + ((u >> 16) & 1u)) >> 16) << 16; + std::memcpy(&v, &u, 4); + return v; +} + +void RouteDump(int64_t layer, const std::vector& logits, + const Kolibri1HostRouting& route, int64_t t, int64_t e, + int64_t top_k) { + if (!StageDumpOn()) return; + for (int64_t i = 0; i < t; ++i) { + double ls = 0; + for (int64_t j = 0; j < e; ++j) + ls += logits[static_cast(i * e + j)]; + std::string ids, ws; + for (int64_t k = 0; k < top_k; ++k) { + const size_t idx = static_cast(i * top_k + k); + ids += (k ? "," : "") + std::to_string(route.ids[idx]); + ws += (k ? "," : "") + + [](float v) { + char b[32]; + std::snprintf(b, sizeof b, "%.6f", v); + return std::string(b); + }(route.weights[idx]); + } + std::fprintf(stderr, + "[STAGE][TT] L%lld route tok=%lld logit_sum=%.6f ids=[%s] " + "w=[%s]\n", + static_cast(layer), static_cast(i), ls, + ids.c_str(), ws.c_str()); + } + if (g_kolibri1_dbg_route_logits.size() == logits.size()) { + double mad = 0, c2 = 0, dd2 = 0; + size_t grid = 0, kmax = 0; + for (size_t k = 0; k < logits.size(); ++k) { + const double c = g_kolibri1_dbg_route_logits[k]; + const double tt = logits[k]; + const double dd = std::fabs(c - tt); + if (dd > mad) { mad = dd; kmax = k; } + c2 += c * c; + dd2 += (c - tt) * (c - tt); + if (Kolibri1DbgToBf16(logits[k]) == logits[k]) ++grid; + } + std::fprintf(stderr, + "[STAGE][TT] L%lld route_diff max=%.6f@%zu cpu=%.4f tt=%.4f " + "bf16grid=%zu/%zu rel_rms=%.5f\n", + static_cast(layer), mad, kmax, + g_kolibri1_dbg_route_logits[kmax], logits[kmax], grid, + logits.size(), std::sqrt(dd2 / (c2 != 0 ? c2 : 1))); + } +} + // ---- The routed-expert refusal (fires by name, counted, never thrown) ------ std::mutex& RefusalMutex() { @@ -309,7 +390,7 @@ DBuf AttentionBlock(Dev d, const Kolibri1AttnWeights& w, const Kolibri1Params& p const auto* f4 = reinterpret_cast(tmp.data()); for (int64_t i = 0; i < tn.Numel(); ++i) s += f4[i]; } - std::fprintf(stderr, "[STAGE] attnin %s sum=%.6f n=%lld dt=%d\n", tag, s, + std::fprintf(stderr, "[STAGE][TT] attnin %s sum=%.6f n=%lld dt=%d\n", tag, s, static_cast(tn.Numel()), static_cast(tn.dtype)); }; @@ -358,7 +439,7 @@ DBuf AttentionBlock(Dev d, const Kolibri1AttnWeights& w, const Kolibri1Params& p }; const int64_t n = t * hq * dh; for (int64_t i = 0; i < n; ++i) s += val(bf[i]); - std::fprintf(stderr, "[STAGE] L%lld pa_out sum=%.6f first=%.6f\n", + std::fprintf(stderr, "[STAGE][TT] L%lld pa_out sum=%.6f first=%.6f\n", static_cast(t * 0), s, val(bf[0])); } return LinearBTDevice(d, o_in, o_w, t); // [T, H] @@ -405,8 +486,67 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, // The inherited sigmoid-logit-add routing (kolibri1_shared.h, verbatim // from the CPU row): selection on logits + bias, weights = sigmoid of // the UNBIASED logits, no renormalisation. + if (StageDumpOn() && l3dbg(layer)) { + const int64_t hk = w.router_gate.shape[1]; + Tensor bt3 = ResidentWeight(d, w.router_gate, {e, hk}); + std::vector tb(static_cast(e) * hk * 2); + d.b.Copy(d.q, tb.data(), bt3.data, tb.size()); + d.b.Synchronize(d.q); + const auto* bf3 = reinterpret_cast(tb.data()); + uint32_t hh = 2166136261u; + double ss = 0; + for (int64_t i = 0; i < e * hk; ++i) { + hh = (hh ^ bf3[i]) * 16777619u; + uint32_t u = static_cast(bf3[i]) << 16; + float f; + std::memcpy(&f, &u, 4); + ss += f; + } + std::fprintf(stderr, "[STAGE][TT] router_gate hash=%08x sum=%.4f\n", hh, ss); + // Host reference logits from the SAME dhn the MoeBlock received and the + // SAME router_gate bytes: separates "the GEMM ate the wrong data" from + // "the GEMM computed badly". + { + const int64_t hk2 = w.router_gate.shape[1]; + std::vector xb(static_cast(t) * hk2 * 2); + d.b.Copy(d.q, xb.data(), dhn.data, xb.size()); + d.b.Synchronize(d.q); + const auto* xb16 = reinterpret_cast(xb.data()); + std::vector href(static_cast(t) * e, 0.f); + for (int64_t n0 = 0; n0 < e; ++n0) { + const auto* wr = + reinterpret_cast(tb.data()) + + static_cast(n0) * hk2; + for (int64_t i = 0; i < t; ++i) { + double acc = 0; + for (int64_t k = 0; k < hk2; ++k) { + auto dec = [](uint16_t b) { + uint32_t u = static_cast(b) << 16; + float f; + std::memcpy(&f, &u, 4); + return static_cast(f); + }; + acc += dec(wr[k]) * dec(xb16[static_cast(i) * hk2 + k]); + } + href[static_cast(i) * e + n0] = static_cast(acc); + } + } + double madd = 0, mlog = 0; + for (size_t k = 0; k < href.size(); ++k) { + madd = std::max(madd, std::fabs(static_cast(href[k]) - + static_cast(logits[k]))); + mlog = std::max(mlog, std::fabs(static_cast(href[k]) - + static_cast( + g_kolibri1_dbg_route_logits[k]))); + } + std::fprintf(stderr, + "[STAGE][TT] router_href maxdiff_dev=%.6f maxdiff_cpu=%.6f\n", + madd, mlog); + } + } Kolibri1HostRouting route = SigmoidLogitAddRouting(logits, bias, t, e, top_k); + RouteDump(layer, logits, route, t, e, top_k); // SLICE ii: the routed tier. Without a streaming context (the // streaming-disabled arm) the B2b-i refusal still fires by name. @@ -490,7 +630,7 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, st.epoch.Record(st.pool.policy); if (StageDumpOn()) { std::fprintf(stderr, - "[STAGE] pool L%lld staged=%lld verified=%lld fills=%lld\n", + "[STAGE][TT] pool L%lld staged=%lld verified=%lld fills=%lld\n", static_cast(layer), static_cast(st.staged_bytes), static_cast(st.readback_verified_bytes), @@ -610,10 +750,18 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, d.b.Copy(d.q, dp + s * rb, sp + static_cast(token_rows[s]) * rb, rb); } - StageDump("moe_gath", layer, gathered, d); + { + char tag[32]; + std::snprintf(tag, sizeof tag, "moe_gath#%d", static_cast(ex)); + StageDump(tag, layer, gathered, d); + } DBuf o = ExpertMlp(d, gathered.t(), ne, inter, it->second.gate, it->second.up, it->second.down); - StageDump("moe_ex", layer, o, d); + { + char tag[32]; + std::snprintf(tag, sizeof tag, "moe_ex#%d", static_cast(ex)); + StageDump(tag, layer, o, d); + } for (int64_t i = 0; i < ne; ++i) { d.b.Copy(d.q, static_cast(expert_out.ptr()) + @@ -630,7 +778,7 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, std::vector wsum(route.weights.begin(), route.weights.end()); double s = 0; for (float v : wsum) s += v; - std::fprintf(stderr, "[STAGE] route sum=%.6f n=%zu\n", s, wsum.size()); + std::fprintf(stderr, "[STAGE][TT] route sum=%.6f n=%zu\n", s, wsum.size()); } Tensor shared_t = shared.t(); vt::MoeCombine(d.q, out.t(), expert_out.t(), dtw.t(), &shared_t, @@ -824,7 +972,7 @@ ForwardLogits ForwardKolibri1TTResidentForward( return f; }; for (int64_t i = 0; i < h; ++i) s += val(bf[i]); - std::fprintf(stderr, "[STAGE] gam sum=%.6f first=[%.6f %.6f %.6f %.6f]\n", + std::fprintf(stderr, "[STAGE][TT] gam sum=%.6f first=[%.6f %.6f %.6f %.6f]\n", s, val(bf[0]), val(bf[1]), val(bf[2]), val(bf[3])); } Tensor dhn_t = dhn.t(); @@ -838,6 +986,40 @@ ForwardLogits ForwardKolibri1TTResidentForward( } StageDump("dhn", l, dhn, d); + if (l == 0 && !g_kolibri1_dbg_dhn.empty() && + g_kolibri1_dbg_dhn.size() == + static_cast(dhn.t().shape[0] * dhn.t().shape[1])) { + std::vector tb(g_kolibri1_dbg_dhn.size()); + d.b.Copy(d.q, tb.data(), dhn.t().data, tb.size() * 2); + d.b.Synchronize(d.q); + auto dec = [](uint16_t b) { + uint32_t u = static_cast(b) << 16; + float f; + std::memcpy(&f, &u, 4); + return static_cast(f); + }; + double d2 = 0; + int worst[3] = {0, 0, 0}; + double wv[3] = {0, 0, 0}; + for (size_t i = 0; i < tb.size(); ++i) { + const double dd = dec(tb[i]) - dec(g_kolibri1_dbg_dhn[i]); + d2 += dd * dd; + if (std::fabs(dd) > wv[2]) { + wv[2] = std::fabs(dd); + worst[2] = static_cast(i); + if (wv[2] > wv[1]) { std::swap(wv[1], wv[2]); std::swap(worst[1], worst[2]); } + if (wv[1] > wv[0]) { std::swap(wv[0], wv[1]); std::swap(worst[0], worst[1]); } + } + } + std::fprintf(stderr, + "[STAGE][TT] dhn_diff rms=%.6f top0 i=%d cpu=%.6f tt=%.6f " + "top1 i=%d cpu=%.6f tt=%.6f top2 i=%d cpu=%.6f tt=%.6f\n", + std::sqrt(d2 / static_cast(tb.size())), worst[0], + dec(g_kolibri1_dbg_dhn[worst[0]]), dec(tb[worst[0]]), + worst[1], dec(g_kolibri1_dbg_dhn[worst[1]]), + dec(tb[worst[1]]), worst[2], + dec(g_kolibri1_dbg_dhn[worst[2]]), dec(tb[worst[2]])); + } // Attention -> post_attn_norm (no residual). DBuf attn = AttentionBlock(d, lw.attn, p, lw.is_sliding, dhn.t(), si.positions.t(), si, *kv_ptr, t, cl.q, cl.k, diff --git a/src/vt/tenstorrent/tenstorrent_ops.cpp b/src/vt/tenstorrent/tenstorrent_ops.cpp index 82af14453a..3ec5e0be55 100644 --- a/src/vt/tenstorrent/tenstorrent_ops.cpp +++ b/src/vt/tenstorrent/tenstorrent_ops.cpp @@ -49,6 +49,19 @@ namespace { // Device compute: keep result on device (CommitDevice2D). Host round-trip only // when the consumer is a host-staged op (EnsureHost) or an untracked buffer. +// The f32-output GEMM contract (the CPU row's LinearBTRaw accumulates in f32 +// and its logits are NOT bf16-representable): request a FLOAT32 output tensor +// AND fp32 destination accumulation. The default config produced a bf16 +// tensor upcast to f32 on commit (every router logit landed on the bf16 grid, +// per-element error up to ~5 against the CPU row) and bf16 partial sums — +// the selection-flip class that drove the L1-moe 3% divergence +// (ISSUE-LOCAL-01M4ER0E9HHM95YZYJB7T5FECN, this branch's record). +std::optional F32AccKernelConfig() { + ttnn::DeviceComputeKernelConfig cfg{}; + cfg.fp32_dest_acc_en = true; + return cfg; +} + void MatmulKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { TT_OP_TRACE("Matmul"); VT_CHECK(a.rank == 2 && b.rank == 2 && out.rank == 2, @@ -68,7 +81,14 @@ void MatmulKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { ttnn::Tensor dev_a = EnsureDevice2D(a, device); ttnn::Tensor dev_b = EnsureMatmulWeightDevice(b, device); if (dev_b.dtype() == ttnn::DataType::BFLOAT8_B) Bfp8MatmulUse(); - ttnn::Tensor dev_c = ttnn::operations::matmul::matmul(dev_a, dev_b); + const std::optional out_dt = + out.dtype == DType::kF32 ? std::optional(ttnn::DataType::FLOAT32) + : std::nullopt; + ttnn::Tensor dev_c = ttnn::operations::matmul::matmul( + dev_a, dev_b, /*transpose_a=*/false, /*transpose_b=*/false, + /*memory_config=*/std::nullopt, out_dt, /*program_config=*/std::nullopt, + /*activation=*/std::nullopt, + out_dt ? F32AccKernelConfig() : std::nullopt); CommitDevice2D(out, std::move(dev_c)); } @@ -77,6 +97,8 @@ void MatmulKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { // exposes a transpose_b flag, so this is the same sequence as kMatmul with // that flag flipped — no separate upload shape needed since `b` is uploaded // in its native [N,K] layout and ttnn transposes on device. + + void MatmulBTKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { TT_OP_TRACE("MatmulBT"); VT_CHECK(a.rank == 2 && b.rank == 2 && out.rank == 2, @@ -96,8 +118,14 @@ void MatmulBTKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { ttnn::Tensor dev_a = EnsureDevice2D(a, device); ttnn::Tensor dev_b = EnsureMatmulWeightDevice(b, device); if (dev_b.dtype() == ttnn::DataType::BFLOAT8_B) Bfp8MatmulUse(); - ttnn::Tensor dev_c = - ttnn::operations::matmul::matmul(dev_a, dev_b, /*transpose_a=*/false, /*transpose_b=*/true); + const std::optional out_dt = + out.dtype == DType::kF32 ? std::optional(ttnn::DataType::FLOAT32) + : std::nullopt; + ttnn::Tensor dev_c = ttnn::operations::matmul::matmul( + dev_a, dev_b, /*transpose_a=*/false, /*transpose_b=*/true, + /*memory_config=*/std::nullopt, out_dt, /*program_config=*/std::nullopt, + /*activation=*/std::nullopt, + out_dt ? F32AccKernelConfig() : std::nullopt); CommitDevice2D(out, std::move(dev_c)); } @@ -461,7 +489,7 @@ void RmsNormKernel(Queue&, Tensor& out, const Tensor& x, const Tensor& weight, sr += LoadElemF32(*residual, j); } std::fprintf(stderr, - "[STAGE] rmsnorm-host call=%d d=%lld sum(x[0..d])=%.6f " + "[STAGE][TT] rmsnorm-host call=%d d=%lld sum(x[0..d])=%.6f " "sum(res[0..d])=%.6f xfirst=[%.6f %.6f] resfirst=[%.6f %.6f]\n", calls.load(), static_cast(d), sx, sr, LoadElemF32(x, 0), LoadElemF32(x, 1), From 0ebb922f6283ec5dec99bb0a18e3a36962009e93 Mon Sep 17 00:00:00 2001 From: Luca Barbato Date: Sun, 11 Oct 2026 07:46:57 +0200 Subject: [PATCH 2/2] =?UTF-8?q?kolibri1=20TT:=20the=20stage-dump=20router?= =?UTF-8?q?=20diagnostic=20indexed=20the=20CPU=20reference=20global=20unco?= =?UTF-8?q?nditionally=20=E2=80=94=20TT-only=20execution=20with=20VT=5FKOL?= =?UTF-8?q?IBRI1=5FTT=5FSTAGE=5FDUMP=3D1=20read=20an=20EMPTY=20process-glo?= =?UTF-8?q?bal=20out=20of=20bounds?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The maintainer review on 1bf1e5c55 was right, and the tree proves it: the b2ii device token gate TT-only with the dump env set SEGFAULTS at layer 0, immediately after the router_gate hash line and before any router_href line (exit 139, /tmp/red-tt-only-prefix.log) — the MoeBlock dump block iterated the host reference and indexed g_kolibri1_dbg_route_logits[k] unconditionally, and the empty vector's null data derefs. A prior CPU run with a SHORTER token batch left the global too small for the same out-of-bounds read, and the size check in the TT RouteDump runs later and cannot protect this earlier loop. The fix separates the two comparisons the block conflated. The HOST-REFERENCE leg (device logits vs host-computed reference logits for the SAME inputs) stands on its own and is bounds-checked over the two local vectors. The CPU leg is now OPTIONAL and runs only through Kolibri1DebugRouteCompare (declared in kolibri1_tt_forward.h, defined out of the TU's anonymous namespace so the host-side tests reach it) when the stashed CPU reference's layer identity AND element size match the current call — a process-global "last CPU layer" cannot establish layer-by-layer parity, so the CPU arm now tags its stash with the layer (g_kolibri1_dbg_route_layer, g_kolibri1_dbg_dhn_layer) and the TT side checks identity before every global read (RouteDump's route_diff and the layer-0 dhn diff included). TT-only, wrong-layer, and short-batch all print cpu_ref=skipped and move on. New device-free cases in test_kolibri1_tt_b2ii.cpp pin the guard: TT-only execution with an empty CPU global (no read, cpu_compared false, the device leg still computed), a shorter prior CPU batch (the exact pre-fix OOB shape), and a different-layer reference (skipped; the same-layer same-size case still compares and reports maxdiff_cpu). On device, post-fix: the TT-only dump run completes the router_href leg on every layer — 50 cpu_ref=skipped lines, no crash. Gate verdicts UNCHANGED (both arms, clean reset per leg): ON 4 HARD flips at the same positions ('der Mond' step 7, 'Australia' step 5, 'Wissen ist Macht' step 8, 'x1 = 3' step 22) then the deterministic memoized-expert DRAM exhaustion abort; host-free OFF 3 HARD flips then the same abort. TT unit suite: only the two known failures (kGdnDecode, kMatmulBTQuantGrouped). FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai/glm-5.3-flash [maki] --- .../models/kolibri1_forward.cpp | 4 + .../models/kolibri1_tt_forward.cpp | 63 +++++++++++--- .../models/kolibri1_tt_forward.h | 37 ++++++++ tests/vllm/models/test_kolibri1_tt_b2ii.cpp | 87 +++++++++++++++++++ 4 files changed, 178 insertions(+), 13 deletions(-) diff --git a/src/vllm/model_executor/models/kolibri1_forward.cpp b/src/vllm/model_executor/models/kolibri1_forward.cpp index 0ac82cc15a..4b0cdafd0f 100644 --- a/src/vllm/model_executor/models/kolibri1_forward.cpp +++ b/src/vllm/model_executor/models/kolibri1_forward.cpp @@ -53,7 +53,9 @@ // SCRATCH DEBUG (layer-1 moe localization): the CPU arm's last router logits, // read by the TT arm's RouteDump for the direct per-expert diff. std::vector g_kolibri1_dbg_route_logits; +int64_t g_kolibri1_dbg_route_layer = -1; std::vector g_kolibri1_dbg_dhn; +int64_t g_kolibri1_dbg_dhn_layer = -1; namespace vllm { @@ -545,6 +547,7 @@ void RouteDump(int64_t layer, const std::vector& logits, ids.c_str(), ws.c_str()); } g_kolibri1_dbg_route_logits = logits; + g_kolibri1_dbg_route_layer = layer; } @@ -650,6 +653,7 @@ ForwardLogits ForwardKolibri1Forward( g_kolibri1_dbg_dhn.assign( reinterpret_cast(dhn.t().data), reinterpret_cast(dhn.t().data) + hn); + g_kolibri1_dbg_dhn_layer = l; } DBuf attn = AttentionBlock(d, lw.attn, p, lw.is_sliding, dhn.t(), si.positions.t(), si, *kv_ptr, t); diff --git a/src/vllm/model_executor/models/kolibri1_tt_forward.cpp b/src/vllm/model_executor/models/kolibri1_tt_forward.cpp index fb6d06f631..d5bedfb954 100644 --- a/src/vllm/model_executor/models/kolibri1_tt_forward.cpp +++ b/src/vllm/model_executor/models/kolibri1_tt_forward.cpp @@ -45,8 +45,8 @@ // SCRATCH DEBUG (layer-1 moe localization): the CPU arm's stashed router // logits (defined in kolibri1_forward.cpp, same process), read by RouteDump. -extern std::vector g_kolibri1_dbg_route_logits; -extern std::vector g_kolibri1_dbg_dhn; +// Declared WITH their layer identity in kolibri1_tt_forward.h; every read +// of these globals is bounds-checked against size AND layer identity. namespace vllm { @@ -174,7 +174,8 @@ void RouteDump(int64_t layer, const std::vector& logits, static_cast(layer), static_cast(i), ls, ids.c_str(), ws.c_str()); } - if (g_kolibri1_dbg_route_logits.size() == logits.size()) { + if (g_kolibri1_dbg_route_layer == layer && + g_kolibri1_dbg_route_logits.size() == logits.size()) { double mad = 0, c2 = 0, dd2 = 0; size_t grid = 0, kmax = 0; for (size_t k = 0; k < logits.size(); ++k) { @@ -532,16 +533,21 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, } } double madd = 0, mlog = 0; - for (size_t k = 0; k < href.size(); ++k) { - madd = std::max(madd, std::fabs(static_cast(href[k]) - - static_cast(logits[k]))); - mlog = std::max(mlog, std::fabs(static_cast(href[k]) - - static_cast( - g_kolibri1_dbg_route_logits[k]))); + const Kolibri1DebugRouteDiff rep = + Kolibri1DebugRouteCompare(layer, logits, href); + madd = rep.maxdiff_dev; + mlog = rep.maxdiff_cpu; + if (rep.cpu_compared) { + std::fprintf( + stderr, + "[STAGE][TT] router_href maxdiff_dev=%.6f maxdiff_cpu=%.6f\n", + madd, mlog); + } else { + std::fprintf(stderr, + "[STAGE][TT] router_href maxdiff_dev=%.6f " + "cpu_ref=skipped(no matching CPU layer/size)\n", + madd); } - std::fprintf(stderr, - "[STAGE][TT] router_href maxdiff_dev=%.6f maxdiff_cpu=%.6f\n", - madd, mlog); } } Kolibri1HostRouting route = @@ -789,6 +795,36 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, } // namespace +Kolibri1DebugRouteDiff Kolibri1DebugRouteCompare( + int64_t layer, const std::vector& logits, + const std::vector& href) { + Kolibri1DebugRouteDiff rep; + // HOST-REFERENCE leg: device logits vs the host-computed reference for + // the SAME inputs. Both vectors are local; bound over the smaller so a + // mismatched caller cannot read either out of range. + const size_t ndev = std::min(logits.size(), href.size()); + for (size_t k = 0; k < ndev; ++k) { + rep.maxdiff_dev = + std::max(rep.maxdiff_dev, std::fabs(static_cast(href[k]) - + static_cast(logits[k]))); + } + // OPTIONAL CPU leg: only against a CPU reference whose layer identity + // AND size match this call. TT-only execution (empty global), a + // different layer, or a shorter prior CPU batch all SKIP cleanly. + if (g_kolibri1_dbg_route_layer == layer && + g_kolibri1_dbg_route_logits.size() == logits.size() && + logits.size() == href.size()) { + rep.cpu_compared = true; + for (size_t k = 0; k < logits.size(); ++k) { + rep.maxdiff_cpu = std::max( + rep.maxdiff_cpu, + std::fabs(static_cast(g_kolibri1_dbg_route_logits[k]) - + static_cast(href[k]))); + } + } + return rep; +} + // ---- The routed-expert refusal (public contract) ----------------------------- std::string Kolibri1TTRoutedExpertRefusalMessage( @@ -986,7 +1022,8 @@ ForwardLogits ForwardKolibri1TTResidentForward( } StageDump("dhn", l, dhn, d); - if (l == 0 && !g_kolibri1_dbg_dhn.empty() && + if (l == 0 && g_kolibri1_dbg_dhn_layer == 0 && + !g_kolibri1_dbg_dhn.empty() && g_kolibri1_dbg_dhn.size() == static_cast(dhn.t().shape[0] * dhn.t().shape[1])) { std::vector tb(g_kolibri1_dbg_dhn.size()); diff --git a/src/vllm/model_executor/models/kolibri1_tt_forward.h b/src/vllm/model_executor/models/kolibri1_tt_forward.h index 7b7e5c5a44..e70c31b46c 100644 --- a/src/vllm/model_executor/models/kolibri1_tt_forward.h +++ b/src/vllm/model_executor/models/kolibri1_tt_forward.h @@ -200,4 +200,41 @@ ForwardLogits ForwardKolibri1TTResidentForward( Kolibri1TTResidentDeviceContext& ctx, Kolibri1TTStreamingDeviceContext* streaming = nullptr); +// ---- SCRATCH DEBUG (layer-1 moe localization): the router diagnostic ------- +// +// The CPU arm stashes its last router logits in +// `g_kolibri1_dbg_route_logits` (kolibri1_forward.cpp) tagged with the +// layer it came from in `g_kolibri1_dbg_route_layer`; the same pattern +// carries the layer-0 dhn stash. The TT arm's diagnostic compares +// against the CPU reference ONLY when the layer identity AND the element +// count match — a process-global "last CPU layer" alone cannot establish +// layer-by-layer parity, and an unconditional index into the global +// reads out of bounds when the TT arm runs alone (the global starts +// empty) or after a shorter CPU batch. +} // namespace vllm + +// Defined at file scope in kolibri1_forward.cpp (global namespace). +extern std::vector g_kolibri1_dbg_route_logits; +extern int64_t g_kolibri1_dbg_route_layer; +extern std::vector g_kolibri1_dbg_dhn; +extern int64_t g_kolibri1_dbg_dhn_layer; + +namespace vllm { + +struct Kolibri1DebugRouteDiff { + double maxdiff_dev = 0; // device logits vs the host reference (same inputs) + double maxdiff_cpu = 0; // device logits vs the CPU arm's stashed reference + bool cpu_compared = false; +}; + +// The MoeBlock router diagnostic: the HOST-REFERENCE leg (device values vs +// host-computed reference logits for the SAME inputs) stands on its own and +// is bounds-checked over the two local vectors. The optional CPU leg runs +// only when the stashed CPU reference exists, its layer identity matches +// `layer`, and its size matches `logits`. Never indexes the CPU global +// without those checks. +Kolibri1DebugRouteDiff Kolibri1DebugRouteCompare( + int64_t layer, const std::vector& logits, + const std::vector& href); + } // namespace vllm diff --git a/tests/vllm/models/test_kolibri1_tt_b2ii.cpp b/tests/vllm/models/test_kolibri1_tt_b2ii.cpp index cb456d8e5a..df9783c81a 100644 --- a/tests/vllm/models/test_kolibri1_tt_b2ii.cpp +++ b/tests/vllm/models/test_kolibri1_tt_b2ii.cpp @@ -1328,3 +1328,90 @@ TEST_CASE("kolibri1 TT B2b-ii: device token gate — the golden argmax chains " // common prefix must be a near-tie. CHECK(instr_hard_flips == 0); } + +// ---- SCRATCH DEBUG: the router diagnostic's reference separation ----------- +// +// RED-FIRST (review finding on 1bf1e5c55): the MoeBlock stage-dump block +// indexed `g_kolibri1_dbg_route_logits[k]` unconditionally while iterating +// the host reference — TT-only execution (the global starts EMPTY) or a +// prior shorter CPU batch read the global OUT OF BOUNDS. The fix separates +// the HOST-REFERENCE leg (same-input TT-vs-host, always safe, both vectors +// local) from the OPTIONAL CPU leg (only when the stashed CPU reference's +// layer identity AND size match this call). These cases pin the guard. + +namespace { +// Saves and restores the process-global CPU debug stashes around each +// case — the globals belong to the CPU arm and the suite must not leak +// state into the device legs. +struct DbgStashGuard { + std::vector route; + int64_t route_layer; + std::vector dhn; + int64_t dhn_layer; + DbgStashGuard() + : route(g_kolibri1_dbg_route_logits), + route_layer(g_kolibri1_dbg_route_layer), + dhn(g_kolibri1_dbg_dhn), + dhn_layer(g_kolibri1_dbg_dhn_layer) {} + ~DbgStashGuard() { + g_kolibri1_dbg_route_logits = route; + g_kolibri1_dbg_route_layer = route_layer; + g_kolibri1_dbg_dhn = dhn; + g_kolibri1_dbg_dhn_layer = dhn_layer; + } +}; +} // namespace + +TEST_CASE("SCRATCH dbg router diagnostic: TT-only execution (empty CPU " + "reference) skips the CPU leg and never reads the global") { + DbgStashGuard guard; + g_kolibri1_dbg_route_logits.clear(); + g_kolibri1_dbg_route_layer = -1; + g_kolibri1_dbg_dhn.clear(); + g_kolibri1_dbg_dhn_layer = -1; + + // TT-only: the dump env is set, the CPU arm never ran in this process. + std::vector logits = {1.0f, -2.0f, 3.0f, 0.5f}; + std::vector href = {1.25f, -2.0f, 2.5f, 0.5f}; + const vllm::Kolibri1DebugRouteDiff rep = + vllm::Kolibri1DebugRouteCompare(0, logits, href); + CHECK_FALSE(rep.cpu_compared); + CHECK(rep.maxdiff_cpu == 0.0); + CHECK(rep.maxdiff_dev == doctest::Approx(0.5)); // |1.0-1.25| vs |3.0-2.5| +} + +TEST_CASE("SCRATCH dbg router diagnostic: a shorter prior CPU batch (size " + "mismatch) skips the CPU leg cleanly") { + DbgStashGuard guard; + // A prior CPU run with a SHORTER token batch left the global too small; + // the pre-fix loop indexed it up to href.size() and read out of bounds. + g_kolibri1_dbg_route_logits = {9.0f, 9.0f}; + g_kolibri1_dbg_route_layer = 0; + + std::vector logits(4, 1.0f); + std::vector href(4, 1.0f); + const vllm::Kolibri1DebugRouteDiff rep = + vllm::Kolibri1DebugRouteCompare(0, logits, href); + CHECK_FALSE(rep.cpu_compared); + CHECK(rep.maxdiff_cpu == 0.0); + CHECK(rep.maxdiff_dev == 0.0); +} + +TEST_CASE("SCRATCH dbg router diagnostic: a different-layer CPU reference " + "never establishes parity") { + DbgStashGuard guard; + g_kolibri1_dbg_route_logits = {9.0f, 9.0f, 9.0f, 9.0f}; + g_kolibri1_dbg_route_layer = 1; // the stash came from layer 1 + + std::vector logits = {0.0f, 0.0f, 0.0f, 0.0f}; + std::vector href = {0.0f, 0.0f, 0.0f, 0.0f}; + const vllm::Kolibri1DebugRouteDiff rep0 = + vllm::Kolibri1DebugRouteCompare(0, logits, href); + CHECK_FALSE(rep0.cpu_compared); // layer identity mismatch: skipped + + const vllm::Kolibri1DebugRouteDiff rep1 = + vllm::Kolibri1DebugRouteCompare(1, logits, href); + CHECK(rep1.cpu_compared); // same layer AND same size: compared + CHECK(rep1.maxdiff_cpu == doctest::Approx(9.0)); + CHECK(rep1.maxdiff_dev == 0.0); +}