kolibri1 TT: fix the residency use-after-free behind the run-to-run nondeterminism

.agents/issues/BACKEND-TENSTORRENT/ISSUE-LOCAL-01M4P21TZ1Q9CA7P66XBH0BQTX.md

@@ -0,0 +1,121 @@
+ID: ISSUE-LOCAL-01M4P21TZ1Q9CA7P66XBH0BQTX
+Title: kolibri1 TT greedy decode is run-to-run nondeterministic on P150 (thalia)
+Row: BACKEND-TENSTORRENT
+State: CLOSED
+Kind: bug
+GitHub: -
+Mirror: PENDING
+Availability: FULL
+Created: 2026-10-11
+Updated: 2026-10-12
+Closed: 2026-10-12
+
+## Problem
+
+On the P150 (host thalia, /tmp/umdtrial-install, /tmp/tt-metal-umdtrial), the kolibri1 Tenstorrent greedy decode chain is not reproducible run to run on identical code and input: three runs of the same one-token-per-step walk produced first tokens 346, 30686, and 10, and the b2ii gate's hard flips move between runs on unmodified origin/main alone. Chain identity is not adjudicable on this path; the b2ii gate is a band gate (141/145, 2.5-nat near-tie teacher-forced). This nondeterminism poisons every measurement. Investigate which decode arm(s) are nondeterministic (host-free ON vs OFF), bisect to the first run-to-run divergent stage via VT_KOLIBRI1_TT_STAGE_DUMP, classify (residency/stale read, async queue race, host thread-pool op ordering, or device FP nondeterminism), and fix if it is a bug in our path.
+
+## Resolution
+
+- ROOT CAUSE (two bugs, both in our path, both on the kolibri1 B2b-ii
+  streaming-MoE memoized dequant cache; classification: residency/stale-read
+  + use-after-free, NOT an async queue race and NOT device FP nondeterminism):
+
+  1. DANGLING-ITERATOR USE-AFTER-FREE in the memo eviction
+     (kolibri1_tt_forward.cpp MoeBlock). The memo is capped at 512 entries;
+     the old order emplaced the new (layer, expert) dequant and THEN evicted
+     the smallest key (`while (size > cap) erase(begin())`). Once the memo's
+     smallest key caught up to the current (layer, expert) — first observed
+     at the 11th forward of the decode bench — the victim WAS the
+     just-emplaced entry: `it` dangled, and the ExpertMlp below read
+     gate/up/down views plus a 7.86 MB dequant buffer that the erase had
+     already be.Free'd. Measured symptom: routed-expert GEMM outputs of
+     uninitialized-DRAM scale (~1e28-1e37) for the caught-up experts while
+     the gathered inputs were bit-identical; the divergence enters the hidden
+     chain at that expert's layer and propagates to the logits.
+
+  2. STALE WeightViewShadow SERVE ON ADDRESS REUSE
+     (tenstorrent_residency.cpp). The memo's up/down expert-weight views are
+     INTERIOR pointers into the 7.86 MB dequant buffer, cached in
+     WeightViewShadows() keyed by the view's host pointer. That map was never
+     invalidated when the buffer was freed: malloc reuses the address for the
+     next dequant (measured: 2357 reused addresses in one 192-step run), and
+     the new expert's up/down GEMMs stale-served the OLD expert's device
+     weights (EnsureWeightViewDevice's hit path returns the cached tensor
+     without comparing bytes). Onset exactly when the memo cap is first
+     crossed and evictions begin (the 3rd forward of the bench walk). The
+     leak also OOM'd the near-full device on the longer b2ii gate (3/3
+     pre-fix runs died at the same 2.5 MB bank allocation; the post-fix gate
+     completes).
+
+  Both arms (host-free ON and OFF) were nondeterministic — the defect is in
+  the shared residency/memo layer, upstream of the arm-specific residual
+  paths. CPU was deterministic (3/3 identical chains). Kolibri1 never
+  captures (no capture call in the TT path), so the only decode modes are the
+  two eager host-free arms.
+
+- BISECTION EVIDENCE: per-step logits/dnorm/hidden/res hashes (scratch,
+  since removed) localized the first divergence to the 3rd forward, in the
+  mid layers, entering the hidden chain through the MoE; the memo trace
+  (FETCH/HIT/MEMO/EVICT) showed the first behavioral difference as the
+  router's expert selection diverging after a corrupted MoE output; the
+  per-expert weight host-byte hashes pinned the first differing weight to a
+  freshly dequanted expert whose memo entry the eviction had just erased
+  (victim == the just-emplaced key). Control experiments: no-eviction runs
+  (raised cap) and leak-the-evicted-entries runs were deterministic through
+  the compared prefix; host-leak-only (device shadow still destroyed) was
+  deterministic, isolating the host free + address reuse as the trigger.
+
+- FIX (row/tt-kolibri-nondet, two changes):
+  1. kolibri1_tt_forward.cpp: evict BEFORE emplacing
+     (Kolibri1TTMemoEvictForEmplace, extracted as a testable helper), so the
+     just-emplaced entry is never the victim and `it` stays valid.
+  2. tenstorrent_residency.cpp: DropWeightViewShadowsInRange sweeps
+     WeightViewShadows + Bfp8WeightShadows for keys inside a retiring
+     allocation, called from UnregisterHostBuffer (be.Free; no drain — the
+     memo eviction's queue is already drained by the preceding scatter
+     readback) and from MarkScratchAcquired (pool re-hand; drains first).
+     Retention was measured and rejected: any retained transient OOMs this
+     near-full device (the prior session's retired-shadow-ring rejection).
+
+- ACCEPTANCE (the deterministic check): 3 identical bench runs per arm,
+  identical full 65-token greedy chains:
+    host-free ON  (default): md5 75f2b5c4a4c12f40410e2a964c333498 x3
+    host-free OFF (VT_TT_HOST_FREE_DECODE=0): md5 89e5b411c6e3790cd090874de36db68a x3
+  Pre-fix the chains differed every run (first tokens 100/44/379 ON,
+  2213/44/861 OFF).
+
+- REGRESSION TESTS: tests/vllm/models/test_kolibri1_tt.cpp — the memo
+  eviction order (the just-emplaced entry survives; the old order erased it
+  — the deterministic red); tests/vt/test_tenstorrent_backend.cpp — freeing
+  a buffer drops the weight-view shadows keyed inside it (the old code leaked
+  the entry — the deterministic red). test_kolibri1_tt: 17/17 pass.
+
+- GATES: TT unit suite (test_tenstorrent_backend): 104/107 — the 3
+  failures are exactly the 2 known failures (kGdnDecode wide-range state,
+  kMatmulBTQuantGrouped P=1 Q4_K) + the 1 known order-state flake (batched
+  decode RAC num_slots=2); no regression, and the new sweep test passes.
+  b2ii gate (both arms): completes deterministically now (pre-fix it OOM'd
+  3/3 this session on the leaked WeightViewShadow tensors, and its flips
+  moved between runs). The gate's CHECKs stay RED by design (the B2b-ii
+  record owns the residual). Post-fix verdicts: ON — ARGMAX CHAIN 0/8
+  (7 hard, 1 near), instrument 140 tf flips (128 HARD, worst 13.5 nats);
+  OFF — 0/8 (7 hard, 1 near), instrument 142 tf flips (128 HARD, worst 16.2
+  nats); identical across repeated runs. All flips are semantic (nat gaps
+  2-16), not garbage.
+
+- FOUND, NOT FIXED (out of scope, reported for the gate's owner): the b2ii
+  ARGMAX CHAIN compares the d-th decoded token against
+  generated_ids[step] instead of generated_ids[step - input_len] — an
+  off-by-input_len index introduced with the one-token prompt walk
+  (9124217b3). The recorded 26/33 baseline predates that rewrite; under the
+  current indexing the compared/match counts are not comparable to it. With
+  the correct index, the post-fix run matches the golden at the first
+  decoded step for 3 of 8 prompts (the rest are semantic near-ties/hard
+  flips, the known B2b-ii residual).
+
+- PERF NOTE: decode ~104 s vs ~73.5 s pre-fix on the bench (+40%): the cost
+  of destroying the ~136k stale weight-view tensors at each memo eviction
+  (correctness requires the drop; retention OOMs the device). Prefill
+  ~214 s vs ~181 s. A generation-counter scheme could cut the destruction
+  count ~30x — follow-up if the decode regression matters.
+

src/vllm/model_executor/models/kolibri1_tt_forward.cpp

@@ -799,7 +799,6 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p,
         *dsts[pi] = MakeTensor(dst, DType::kBF16, d.q.device,
                                std::vector<int64_t>{n, k});
       }
-      it = st.memo.emplace(key, std::move(dq)).first;
       // THE BOUND: the memo is a cache, and an unbounded one exhausts the
       // device (each entry owns a 3*inter*h*2 = 7.86 MB bf16 buffer at the
       // Kolibri-1 dims; a free-running decode routes to ever-new experts
@@ -808,8 +807,18 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p,
       // resident context and the slot pool. Eviction is correctness-free
       // (a dropped entry is re-dequanted from the authoritative slot bytes
       // on the next touch), and a fresh dequant is bit-identical.
+      //
+      // EVICT BEFORE EMPLACING. The old order (emplace, then erase begin())
+      // erased the JUST-EMPLACED entry once the memo's smallest key caught
+      // up to the current (layer, expert) — from that step on, `it` dangled
+      // and the ExpertMlp below read gate/up/down views and a dequant buffer
+      // that had already been freed (be.Free inside the erase): garbage
+      // weights, uninitialized-DRAM-scale outputs, and run-to-run
+      // nondeterministic decode. Evicting first keeps the current entry out
+      // of the victim set, so `it` stays valid for the GEMM below.
       constexpr size_t kMemoMaxEntries = 512;
-      while (st.memo.size() > kMemoMaxEntries) st.memo.erase(st.memo.begin());
+      Kolibri1TTMemoEvictForEmplace(st.memo, kMemoMaxEntries);
+      it = st.memo.emplace(key, std::move(dq)).first;
     }
     const int64_t ne = static_cast<int64_t>(token_rows.size());
     DBuf gathered(d, DType::kBF16, {ne, h});
@@ -890,6 +899,14 @@ Kolibri1DebugRouteDiff Kolibri1DebugRouteCompare(
   return rep;
 }
 
+// ---- The memo eviction (public contract) -------------------------------------
+//
+// See the header: evict BEFORE the caller's emplace so the just-emplaced
+// entry is never the victim.
+void Kolibri1TTMemoEvictForEmplace(Kolibri1TTMemo& memo, size_t cap) {
+  while (memo.size() >= cap) memo.erase(memo.begin());
+}
+
 // ---- The routed-expert refusal (public contract) -----------------------------
 
 std::string Kolibri1TTRoutedExpertRefusalMessage(
@@ -1007,6 +1024,7 @@ ForwardLogits ForwardKolibri1TTResidentForward(
                  static_cast<long long>(p.num_experts_per_tok));
   }
 
+
   // Embedding (bf16 table, [vocab, H] raw orientation) — the shared
   // residency seam, exactly like the CPU row.
   std::shared_ptr<void> hidden_hold;

src/vllm/model_executor/models/kolibri1_tt_forward.h

@@ -160,6 +160,22 @@ struct Kolibri1TTStreamingDeviceContext {
   std::map<std::pair<int64_t, int64_t>, SlotDequant> memo;
 };
 
+// The memo container type, named for the eviction helper below.
+using Kolibri1TTMemo =
+    std::map<std::pair<int64_t, int64_t>,
+             Kolibri1TTStreamingDeviceContext::SlotDequant>;
+
+// Evicts `memo` down to `cap` - 1 entries, smallest key first, so the
+// caller's IMMEDIATELY FOLLOWING emplace lands at size <= cap and is never
+// the eviction victim. The caller must emplace right after this returns and
+// must not hold iterators across the call. Evicting AFTER the emplace (the
+// old order) erased the just-emplaced entry once the memo's smallest key
+// caught up to the current (layer, expert) — a dangling-iterator
+// use-after-free that fed freed dequant buffers to the expert GEMMs and
+// made the decode run-to-run nondeterministic (host-side testable, no
+// card: tests/vllm/models/test_kolibri1_tt.cpp).
+void Kolibri1TTMemoEvictForEmplace(Kolibri1TTMemo& memo, size_t cap);
+
 // Builds the streaming context: the B1 streaming plan over the REAL
 // checkpoint's byte shape (device_budget_bytes, kv_reserve_bytes from the
 // options), the slot pool plan, and the per-layer device pool allocations.

src/vt/tenstorrent/tenstorrent_device.h

@@ -225,6 +225,12 @@ int64_t DeviceDramTotalBytes();
 // EnsureDevice2D does (creating the slot's persistent buffer), reproducing
 // the loader's weight staging before a keep-quant matmul.
 void StageWeightBf16ForTest(const Tensor& t, MeshDevice& device);
+// Stale-serve-fix test hook: stage `t` exactly as EnsureMatmulWeightDevice
+// routes an interior bf16 weight view (a WeightViewShadows entry keyed by
+// the view's host pointer), and report the entry count, so a focused case
+// can verify that freeing the owning buffer drops the entries inside it.
+void StageWeightViewForTest(const Tensor& t);
+size_t WeightViewShadowCountForTest();
 // W4d W6: stage the keep-quant word shadow for a block-quant weight at
 // LOAD time (the residency pre-pass; no-op for non-block dtypes).
 void StageKeepQuantWordsFor(const Tensor& packed);
@@ -280,6 +286,8 @@ inline void ResetAllocTraceForTest() {}
 inline int64_t FreeDeviceDramBytesForTest() { return 0; }
 inline int64_t DeviceDramTotalBytes() { return 0; }
 inline void StageWeightBf16ForTest(const Tensor&, MeshDevice&) {}
+inline void StageWeightViewForTest(const Tensor&) {}
+inline size_t WeightViewShadowCountForTest() { return 0; }
 inline void StageKeepQuantWordsFor(const Tensor&) {}
 inline void DumpSlotCensus(const char*) {}
 inline void ReleaseWarmShapeSlots(uint32_t) {}

src/vt/tenstorrent/tenstorrent_residency.cpp

@@ -646,6 +646,52 @@ std::map<uintptr_t, Bfp8WeightShadow>& Bfp8WeightShadows() {
   return *m;
 }
 
+// Invalidate every pointer-keyed weight-view shadow whose key lies within
+// the retiring allocation [host, host+bytes). These maps are keyed by the
+// VIEW's host pointer — interior offsets included — so an entry outlives
+// the buffer it was staged from unless it is dropped here. The address is
+// about to be reused (a pool re-hand, or a free whose malloc successor lands
+// on the same address), and a surviving entry would serve the PREVIOUS
+// tenant's device bytes to the new tenant's views. That is the run-to-run
+// decode nondeterminism this closes: the kolibri1 memo eviction frees 7.86 MB
+// dequant buffers whose up/down expert-weight views are interior pointers
+// cached in WeightViewShadows; malloc reuses the address for the next
+// dequant, and the new expert's GEMM stale-served the OLD expert's weights
+// (onset exactly when the memo cap is first crossed and evictions begin).
+// The dropped tensors' consumers may still be queued when `drain` is set
+// (a pool re-hand is not a drain point); the be.Free path passes false —
+// its mid-run caller is the kolibri1 memo eviction, whose queue was already
+// drained by the preceding expert's scatter readback.
+void DropWeightViewShadowsInRange(void* host, size_t bytes, bool drain) {
+  if (host == nullptr || bytes == 0) return;
+  const uintptr_t lo = reinterpret_cast<uintptr_t>(host);
+  const uintptr_t hi = lo + bytes;
+  std::vector<uintptr_t> wv_dead, bfp8_dead;
+  {
+    std::lock_guard<std::mutex> g(WeightViewMutex());
+    for (auto it = WeightViewShadows().lower_bound(lo);
+         it != WeightViewShadows().end() && it->first < hi; ++it)
+      wv_dead.push_back(it->first);
+  }
+  {
+    std::lock_guard<std::mutex> g(Bfp8WeightMutex());
+    for (auto it = Bfp8WeightShadows().lower_bound(lo);
+         it != Bfp8WeightShadows().end() && it->first < hi; ++it)
+      bfp8_dead.push_back(it->first);
+  }
+  if (wv_dead.empty() && bfp8_dead.empty()) return;
+  if (drain && !tt_capture_active())
+    SharedMeshDevice().mesh_command_queue().finish();
+  {
+    std::lock_guard<std::mutex> g(WeightViewMutex());
+    for (const uintptr_t k : wv_dead) WeightViewShadows().erase(k);
+  }
+  {
+    std::lock_guard<std::mutex> g(Bfp8WeightMutex());
+    for (const uintptr_t k : bfp8_dead) Bfp8WeightShadows().erase(k);
+  }
+}
+
 ttnn::Tensor EnsureBfp8WeightDevice(const Tensor& b, MeshDevice& device) {
   const uint32_t rows = static_cast<uint32_t>(b.shape[0]);
   const uint32_t cols = static_cast<uint32_t>(b.shape[1]);
@@ -875,6 +921,20 @@ ttnn::Tensor EnsureDevice2D(const Tensor& t, MeshDevice& device) {
 // needs external linkage for the focused test and calls EnsureDevice2D,
 // which stays TU-internal here.)
 
+// Test surface for the stale-serve fix (external linkage, same pattern):
+// StageWeightViewForTest stages an INTERIOR bf16 weight view exactly as
+// EnsureMatmulWeightDevice routes it (a WeightViewShadows entry keyed by
+// the view's pointer), and WeightViewShadowCountForTest reports the map
+// size, so a focused case can verify that freeing the owning buffer drops
+// the entries keyed inside it.
+void StageWeightViewForTest(const Tensor& t) {
+  (void)EnsureWeightViewDevice(t, SharedMeshDevice());
+}
+size_t WeightViewShadowCountForTest() {
+  std::lock_guard<std::mutex> g(WeightViewMutex());
+  return WeightViewShadows().size();
+}
+
 // DEBUG (BACKEND-TENSTORRENT-QWEN35 W2c): ensure the tensor is staged on
 // device exactly as a consuming kernel would, then read the DEVICE copy back.
 // Comparing this against the host master exposes staging corruption that a
@@ -1047,12 +1107,23 @@ void RegisterHostBuffer(void* host, size_t bytes) {
 
 void UnregisterHostBuffer(void* host) {
   if (host == nullptr) return;
+  size_t bytes = 0;
   {
     std::lock_guard<std::mutex> g(SlotMutex());
     if (std::getenv("VT_TT_SLOT_TRACE") != nullptr)
       std::fprintf(stderr, "[TT-SLOT] unregister %p\n", host);
+    // Immediate destroy (no drain): be.Free's mid-run caller is the kolibri1
+    // memo eviction, whose 7.86 MB shadow frees the very next GEMM
+    // allocations reuse, and whose slot's consumers are already complete (the
+    // preceding expert's scatter readback drained the queue).
+    auto it = Slots().find(reinterpret_cast<uintptr_t>(host));
+    if (it != Slots().end()) bytes = it->second.bytes;
     Slots().erase(reinterpret_cast<uintptr_t>(host));
   }
+  // The freed address is about to be reused: drop the pointer-keyed
+  // weight-view shadows keyed inside it (stale-serve fix). No drain: the
+  // queue was already drained past these shadows' consumers.
+  DropWeightViewShadowsInRange(host, bytes, /*drain=*/false);
   DropPagedKvShadow(host);
   DropEmbedTableShadow(host);
   DropDecodedWeightShadow(host);
@@ -1082,9 +1153,11 @@ void MarkHostWritten(void* host) {
 
 void MarkScratchAcquired(void* host) {
   if (host == nullptr) return;
+  size_t bytes = 0;
   {
     std::lock_guard<std::mutex> g(SlotMutex());
     BufferSlot* s = FindSlot(host);
+    if (s != nullptr) bytes = s->bytes;
     if (s != nullptr) {
       // The state MarkHostWritten registers, plus the W7 reservation: the
       // bytes on BOTH sides belong to the previous tenant. host_current stays
@@ -1100,6 +1173,10 @@ void MarkScratchAcquired(void* host) {
   // Same drop MarkHostWritten performed for every caller — a pool block must
   // never serve a stale embed-table staging either.
   DropEmbedTableShadow(host);
+  // The re-handed address keeps its bytes: drop the pointer-keyed
+  // weight-view shadows keyed inside it (stale-serve fix). A pool re-hand
+  // is not a drain point, so the queue is drained first.
+  DropWeightViewShadowsInRange(host, bytes, /*drain=*/true);
 }
 
 // TRUSTED dump (BACKEND-TENSTORRENT-QWEN35 W2c measurement reset): whole