From f076ce92cc43d6053ee9187013a727a13b81fd2f Mon Sep 17 00:00:00 2001 From: Luca Barbato Date: Mon, 12 Oct 2026 06:56:24 +0200 Subject: [PATCH] =?UTF-8?q?fix(BACKEND-TENSTORRENT):=20make=20the=20kolibr?= =?UTF-8?q?i1=20TT=20greedy=20decode=20deterministic=20=E2=80=94=20two=20m?= =?UTF-8?q?emo/residency=20use-after-free=20bugs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kolibri1 TT greedy decode chain was not reproducible run to run on identical code and input (first tokens 100/44/379 across three runs on the host-free ON arm, 2213/44/861 on OFF), poisoning every measurement on the P150. CPU was deterministic (3/3 identical chains); kolibri1 never captures, so both eager host-free arms shared the defect. Root cause: two bugs in the B2b-ii streaming-MoE memoized dequant cache, both in our path (not an async queue race, not device FP nondeterminism). 1. Dangling-iterator use-after-free in the memo eviction (kolibri1_tt_forward.cpp MoeBlock). The memo caps at 512 entries and the old order emplaced the new (layer, expert) dequant and THEN evicted the smallest key. Once the memo's smallest key caught up to the current (layer, expert) — first at the 11th forward of the decode bench — the victim was the just-emplaced entry: `it` dangled and the ExpertMlp read gate/up/down views plus a 7.86 MB dequant buffer the erase had already be.Free'd. Measured: routed-expert GEMM outputs of uninitialized-DRAM scale (~1e28-1e37) with bit-identical gathered inputs. Fixed by evicting BEFORE the emplace (Kolibri1TTMemoEvictForEmplace, extracted as a testable helper) so the current entry is never the victim. 2. Stale WeightViewShadow serve on address reuse (tenstorrent_residency.cpp). The memo's up/down expert-weight views are interior pointers into the freed dequant buffer, cached in WeightViewShadows keyed by the view's host pointer; the map was never invalidated on free. malloc reuses the address for the next dequant (2357 reused addresses in one 192-step run) and the new expert's GEMMs stale-served the OLD expert's weights. Onset exactly when the memo cap is first crossed (the 3rd forward). The leak also OOM'd the near-full device on the b2ii gate (3/3 pre-fix runs died at the same 2.5 MB bank allocation). Fixed by sweeping WeightViewShadows + Bfp8WeightShadows for keys inside a retiring allocation, at UnregisterHostBuffer (no drain — the eviction's queue is already drained by the preceding scatter readback) and at MarkScratchAcquired (pool re-hand; drains first). Retention was measured and rejected: any retained transient OOMs this device. Acceptance (the deterministic check): 3 identical bench runs per arm, identical full 65-token greedy chains — ON md5 75f2b5c4a4c12f40410e2a964c333498 x3, OFF md5 89e5b411c6e3790cd090874de36db68a x3. Regression tests: the memo eviction order (test_kolibri1_tt.cpp, 17/17 pass) and the weight-view sweep (test_tenstorrent_backend.cpp; the suite stays at the 2 known failures + the 1 known order-state flake, no regression). b2ii gate completes deterministically on both arms now (pre-fix it OOM'd or moved); its CHECKs stay RED by design (the B2b-ii record owns the residual) with semantic flips, not garbage. Found but not fixed (reported, out of scope): the b2ii ARGMAX CHAIN compares generated_ids[step] instead of generated_ids[step - input_len] — an off-by-input_len index from the one-token prompt walk (9124217b3), so its counts are not comparable to the pre-walk 26/33 baseline. Issue: ISSUE-LOCAL-01M4P21TZ1Q9CA7P66XBH0BQTX FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai/glm-5.3-flash [maki] --- .../ISSUE-LOCAL-01M4P21TZ1Q9CA7P66XBH0BQTX.md | 121 ++++++++++++++++++ .../models/kolibri1_tt_forward.cpp | 22 +++- .../models/kolibri1_tt_forward.h | 16 +++ src/vt/tenstorrent/tenstorrent_device.h | 8 ++ src/vt/tenstorrent/tenstorrent_residency.cpp | 77 +++++++++++ tests/vllm/models/test_kolibri1_tt.cpp | 53 ++++++++ tests/vt/test_tenstorrent_backend.cpp | 36 ++++++ 7 files changed, 331 insertions(+), 2 deletions(-) create mode 100644 .agents/issues/BACKEND-TENSTORRENT/ISSUE-LOCAL-01M4P21TZ1Q9CA7P66XBH0BQTX.md diff --git a/.agents/issues/BACKEND-TENSTORRENT/ISSUE-LOCAL-01M4P21TZ1Q9CA7P66XBH0BQTX.md b/.agents/issues/BACKEND-TENSTORRENT/ISSUE-LOCAL-01M4P21TZ1Q9CA7P66XBH0BQTX.md new file mode 100644 index 0000000000..170137b3b2 --- /dev/null +++ b/.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. + diff --git a/src/vllm/model_executor/models/kolibri1_tt_forward.cpp b/src/vllm/model_executor/models/kolibri1_tt_forward.cpp index 2e5d68a30e..9d1e9bf794 100644 --- a/src/vllm/model_executor/models/kolibri1_tt_forward.cpp +++ b/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{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(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(p.num_experts_per_tok)); } + // Embedding (bf16 table, [vocab, H] raw orientation) — the shared // residency seam, exactly like the CPU row. std::shared_ptr hidden_hold; diff --git a/src/vllm/model_executor/models/kolibri1_tt_forward.h b/src/vllm/model_executor/models/kolibri1_tt_forward.h index e70c31b46c..53995c1a69 100644 --- a/src/vllm/model_executor/models/kolibri1_tt_forward.h +++ b/src/vllm/model_executor/models/kolibri1_tt_forward.h @@ -160,6 +160,22 @@ struct Kolibri1TTStreamingDeviceContext { std::map, SlotDequant> memo; }; +// The memo container type, named for the eviction helper below. +using Kolibri1TTMemo = + std::map, + 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. diff --git a/src/vt/tenstorrent/tenstorrent_device.h b/src/vt/tenstorrent/tenstorrent_device.h index 7eb633ea60..f438dd3998 100644 --- a/src/vt/tenstorrent/tenstorrent_device.h +++ b/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) {} diff --git a/src/vt/tenstorrent/tenstorrent_residency.cpp b/src/vt/tenstorrent/tenstorrent_residency.cpp index 131ad1c775..fe63b189b6 100644 --- a/src/vt/tenstorrent/tenstorrent_residency.cpp +++ b/src/vt/tenstorrent/tenstorrent_residency.cpp @@ -646,6 +646,52 @@ std::map& 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(host); + const uintptr_t hi = lo + bytes; + std::vector wv_dead, bfp8_dead; + { + std::lock_guard 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 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 g(WeightViewMutex()); + for (const uintptr_t k : wv_dead) WeightViewShadows().erase(k); + } + { + std::lock_guard 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(b.shape[0]); const uint32_t cols = static_cast(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 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 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(host)); + if (it != Slots().end()) bytes = it->second.bytes; Slots().erase(reinterpret_cast(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 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 diff --git a/tests/vllm/models/test_kolibri1_tt.cpp b/tests/vllm/models/test_kolibri1_tt.cpp index cec594c6e5..166a07e84e 100644 --- a/tests/vllm/models/test_kolibri1_tt.cpp +++ b/tests/vllm/models/test_kolibri1_tt.cpp @@ -23,6 +23,7 @@ #include #include "vllm/model_executor/models/kolibri1_tt.h" +#include "vllm/model_executor/models/kolibri1_tt_forward.h" #include "vllm/model_executor/models/kolibri1_weights.h" #include "vllm/model_executor/model_loader/safetensors_reader.h" #include "vllm/model_executor/models/model_registry.h" @@ -788,3 +789,55 @@ TEST_CASE("kolibri1 TT reset lane: the slot-table fingerprint flips on a " const uint64_t after1 = policy.Fingerprint(); CHECK_FALSE(policy.ContentChangedSince(after1)); } + +TEST_CASE("kolibri1 TT memo eviction: the just-emplaced entry is never the " + "victim (dangling-iterator use-after-free regression)") { + // The memo cap is 512 on the production shape and the eviction erases the + // smallest key. 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): the iterator the expert loop holds dangled, + // and the ExpertMlp read gate/up/down views plus 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. Kolibri1TTMemoEvictForEmplace evicts BEFORE the caller's + // emplace so the current entry is never the victim. + constexpr size_t kCap = 512; + using SlotDequant = Kolibri1TTStreamingDeviceContext::SlotDequant; + Kolibri1TTMemo memo; + + // Fill to the cap with keys (0,1)..(0,cap): the smallest key is (0,1), + // the catch-up shape the bug needed (a fresh dequant whose key sorts + // below every resident entry). + for (int64_t e = 1; e <= static_cast(kCap); ++e) + memo.emplace(std::make_pair(int64_t{0}, e), SlotDequant{}); + REQUIRE(memo.size() == kCap); + + // The caller's next dequant is for the smallest key (0,0). The old + // emplace-then-evict order erased (0,0) itself here; the helper evicts + // the pre-existing smallest (0,1) instead, and the emplace that follows + // lands at the cap with a valid iterator. + Kolibri1TTMemoEvictForEmplace(memo, kCap); + const auto key = std::make_pair(int64_t{0}, int64_t{0}); + auto it = memo.emplace(key, SlotDequant{}).first; + + CHECK(memo.size() == kCap); + // The just-emplaced entry SURVIVES (the old order erased it — this is the + // deterministic red on the old code). + CHECK(memo.find(key) != memo.end()); + // The iterator the expert loop holds stays valid and names the entry. + CHECK(it != memo.end()); + CHECK(it->first == key); + // The victim was a PRE-EXISTING entry: (0,1) is gone, (0,0) is resident. + CHECK(memo.find(std::make_pair(int64_t{0}, int64_t{1})) == memo.end()); + CHECK(memo.begin()->first == key); + + // A second round on the full memo keeps the invariant: emplacing the + // current smallest key never destroys it either. + Kolibri1TTMemoEvictForEmplace(memo, kCap); + const auto key2 = std::make_pair(int64_t{0}, int64_t{1}); + auto it2 = memo.emplace(key2, SlotDequant{}).first; + CHECK(memo.size() == kCap); + CHECK(memo.find(key2) != memo.end()); + CHECK(it2 != memo.end()); + CHECK(it2->first == key2); +} diff --git a/tests/vt/test_tenstorrent_backend.cpp b/tests/vt/test_tenstorrent_backend.cpp index eae32f09c7..b88b804625 100644 --- a/tests/vt/test_tenstorrent_backend.cpp +++ b/tests/vt/test_tenstorrent_backend.cpp @@ -12489,3 +12489,39 @@ TEST_CASE("kTENSTORRENT kSigmoidBiasTopK matches the CPU reference " CHECK(ci[1] == 2); } } +TEST_CASE("tt residency: freeing a buffer drops the weight-view shadows " + "keyed inside it (stale-serve regression)") { + if (!TenstorrentPresent()) { + MESSAGE("SKIPPED: no Tenstorrent device on this box"); + return; + } + Backend& backend = vt::GetBackend(DeviceType::kTENSTORRENT); + // A 64 KiB host allocation; the weight view is an INTERIOR bf16 window — + // exactly the kolibri1 memo dequant's up/down view shape (an interior + // pointer cached in WeightViewShadows, keyed by the view's host pointer). + constexpr size_t kBytes = 64 * 1024; + void* host = backend.Alloc(kBytes); + std::vector pattern(kBytes / 2); + for (size_t i = 0; i < pattern.size(); ++i) + pattern[i] = static_cast(i % 251); + std::memcpy(host, pattern.data(), kBytes); + + auto* view_base = static_cast(host) + kBytes / 2; + Tensor view = Tensor::Contiguous(view_base, vt::DType::kBF16, + Device{DeviceType::kTENSTORRENT, 0}, + {32, 32}); + const size_t before = vt::tenstorrent::WeightViewShadowCountForTest(); + vt::tenstorrent::StageWeightViewForTest(view); + // The interior view staged its own persistent shadow, keyed by its + // pointer (NOT the base's slot — the interior-view rule). + CHECK(vt::tenstorrent::WeightViewShadowCountForTest() == before + 1); + + // Freeing the owning buffer must drop the entry keyed inside it: the + // address is about to be reused, and a surviving entry would serve the + // previous tenant's device bytes to the new tenant's view (the old code + // leaked the entry — the kolibri1 memo eviction's address reuse then + // stale-served the OLD expert's weights to the new expert's GEMM). + backend.Free(host); + CHECK(vt::tenstorrent::WeightViewShadowCountForTest() == before); +} +