kolibri1 TT: the production-shape decode throughput bench

.agents/issues/MODEL-TEXT-kolibri-1-tenstorrent/ISSUE-LOCAL-01M4MHZM17MAAJJF1Z071YA7R0.md

@@ -0,0 +1,28 @@
+ID: ISSUE-LOCAL-01M4MHZM17MAAJJF1Z071YA7R0
+Title: kolibri1 TT streaming MoE: a miss fetch's LRU eviction can evict an expert the same step selected
+Row: MODEL-TEXT-kolibri-1-tenstorrent
+State: CLOSED
+Kind: bug
+GitHub: -
+Mirror: PENDING
+Availability: FULL
+Created: 2026-10-11
+Updated: 2026-10-11
+Closed: 2026-10-11
+
+## Problem
+
+In MoeBlock's routed path (src/vllm/model_executor/models/kolibri1_tt_forward.cpp), PlanKolibri1TTMoEDispatch classifies residents/misses against the pre-fetch table, then the fetch loop Touch()es each miss. Touch on a full table evicts the global LRU slot, which can BE an expert this step's route selected as a resident; the compute loop then hits the 'expert N selected on layer L but no slot holds it' miss-handle refusal on a fully warm table. Observed twice on the production-shape decode bench (row/tt-kolibri-decode-bench), model /mnt/models/Aleph-Alpha/Kolibri-1, gfx under the GPU lease: layer 3, expert 258 resident per dispatch, evicted by the same step's expert-323 fetch; earlier run evicted 374. Resident slots' recency is only stamped at fill time, never on reuse, so a long-resident but-this-step-needed expert is the prime victim.
+
+## Resolution
+
+2026-10-11: FIXED in the same flow (branch row/tt-kolibri-decode-bench,
+src/vllm/model_executor/models/kolibri1_tt_forward.cpp, the MoeBlock routed
+path): the fetch loop now refreshes every SAME-STEP resident's recency
+(Touch on a resident only restamps it) BEFORE staging the misses, so a
+miss's LRU victim can no longer be an expert this step's route selected.
+Evidence: the bench previously refused deterministically at layer 3
+(experts 374 then 258 across runs); with the fix the bench completes a
+192-forward free run on the card in BOTH arms (host-free on and off), and
+the b2ii device token gate stays 13/13, b2bi 10/10, tt 16/16.
+

.agents/issues/MODEL-TEXT-kolibri-1-tenstorrent/ISSUE-LOCAL-01M4MJG2CKSM1FK53PZ44QFN1J.md

@@ -0,0 +1,28 @@
+ID: ISSUE-LOCAL-01M4MJG2CKSM1FK53PZ44QFN1J
+Title: kolibri1 TT streaming MoE: the slot-dequant memo grows without bound and exhausts device DRAM on free-running decode
+Row: MODEL-TEXT-kolibri-1-tenstorrent
+State: CLOSED
+Kind: bug
+GitHub: -
+Mirror: PENDING
+Availability: FULL
+Created: 2026-10-11
+Updated: 2026-10-11
+Closed: 2026-10-11
+
+## Problem
+
+Kolibri1TTStreamingDeviceContext::memo (src/vllm/model_executor/models/kolibri1_tt_forward.h:160) caches one bf16 dequant per (layer, expert) ever routed to, 3*inter*h*2 = 7.86 MB per entry at the Kolibri-1 dims (h=2560, inter=512, 50 layers, 384 experts, top-6). Nothing bounds the map: a free-running decode selects ever-new experts and the entries' device buffers accumulate until the tt-metal allocator refuses (observed on the production-shape decode bench: TT_FATAL Out of Memory, 2.6 MB request with 300 KB free, deterministic at prefill step ~75 of the first run; the b2ii/b2bi gates never hit it because they run tens of forwards). The reset lane only clears on a slot-content change, which is not a growth bound.
+
+## Resolution
+
+2026-10-11: FIXED in the same flow (branch row/tt-kolibri-decode-bench,
+src/vllm/model_executor/models/kolibri1_tt_forward.cpp, the memo insert):
+the memo is capped at 512 entries (~4 GiB bf16, beside the 3.5 GiB resident
+context and the slot pool). Eviction is correctness-free: a dropped entry is
+re-dequanted bit-identically from the authoritative, readback-verified slot
+bytes on the next touch. Evidence: without the cap the bench died with
+TT_FATAL Out of Memory (2.6 MB request, 300 KB free) deterministically at
+prefill step ~75; with the cap it completes a 192-forward free run in both
+arms, and the b2ii/b2bi/tt gates stay green (13/13, 10/10, 16/16).
+

src/vllm/model_executor/models/kolibri1_tt_forward.cpp

@@ -585,6 +585,15 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p,
   vt::Backend& be = vt::GetBackend(d.q.device.type);
   std::vector<uint8_t> slot_image(
       static_cast<size_t>(st.pool.packed_bytes_per_slot));
+  // The same-step residents must not become eviction victims: refresh
+  // their recency BEFORE the fetch loop, so a miss's LRU victim is never
+  // an expert THIS step's route selected (observed on the production-shape
+  // decode bench: expert 258 resident on layer 3, then evicted by the same
+  // step's 323 fetch, and the compute loop refused "no slot holds it" — a
+  // self-inflicted miss-handle on a fully warm table). Touch on a resident
+  // only refreshes the stamp; the table is unchanged.
+  for (const int64_t e : ld.resident_experts)
+    st.pool.policy.Touch(layer, e);
   for (const Kolibri1TTSlotFetchJob& job : list.jobs) {
     // Touch FIRST: the policy's slot is authoritative (the dispatch
     // plan's fetch_slots are advisory — SuggestSlotFor is read-only over
@@ -745,6 +754,16 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p,
                                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
+      // until the allocator refuses — ISSUE-LOCAL-01M4MJG2CKSM1FK53PZ44QFN1J).
+      // Cap it: 512 entries ≈ 4 GiB, comfortably beside the 3.5 GiB
+      // 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.
+      constexpr size_t kMemoMaxEntries = 512;
+      while (st.memo.size() > kMemoMaxEntries) st.memo.erase(st.memo.begin());
     }
     const int64_t ne = static_cast<int64_t>(token_rows.size());
     DBuf gathered(d, DType::kBF16, {ne, h});

tests/CMakeLists.txt

@@ -230,6 +230,16 @@ if(VLLM_CPP_TENSTORRENT)
   target_include_directories(test_kolibri1_tt_b2ii PRIVATE ${CMAKE_SOURCE_DIR}/src)
   target_compile_definitions(test_kolibri1_tt_b2ii PRIVATE
     KOLIBRI1_GOLDENS="${CMAKE_CURRENT_SOURCE_DIR}/vllm/models/kolibri1_goldens.json")
+  # test_kolibri1_tt_decode_bench is the production-shape decode THROUGHPUT
+  # bench on the TT device path (the TT mirror of the CPU
+  # test_kolibri1_decode_bench): load once, one 128-token prefill, 64
+  # free-running greedy steps through ModelRegistry::Forward, wall/prefill/
+  # decode split and the greedy chain PRINTED, never asserted. Operator-run
+  # under the GPU lease (VT_KOLIBRI1_TT_B2II_MODEL gate).
+  vllm_cpp_add_test(test_kolibri1_tt_decode_bench
+                    vllm/models/test_kolibri1_tt_decode_bench.cpp)
+  target_include_directories(test_kolibri1_tt_decode_bench
+                             PRIVATE ${CMAKE_SOURCE_DIR}/src)
 endif()
 # mimo_v2_weights.h is a MODEL-PRIVATE header under src/, not include/vllm/.
 vllm_cpp_add_test(test_mimov2_w2 vllm/models/test_mimov2_w2.cpp)

tests/vllm/models/test_kolibri1_tt_decode_bench.cpp

@@ -0,0 +1,280 @@
+// Kolibri-1 Tenstorrent decode throughput bench (production shape)
+// (MODEL-TEXT-kolibri-1-tenstorrent; spec .agents/specs/kolibri-tt.md).
+//
+// This is the TT-device mirror of tests/vllm/models/
+// test_kolibri1_decode_bench.cpp (the CPU production-shape gate): load once,
+// ONE 128-token prefill, then 64 free-running greedy decode steps through the
+// production registry seam (ModelRegistry::Forward, pure_decode), timed as
+// wall / prefill / decode with the full greedy chain printed. NO per-op
+// dumps, NO per-step CPU comparison, NO teacher forcing, NO token anchor —
+// the numerics adjudication is a separate question; this TU measures.
+//
+// Skip convention (the b2bi device-leg gate): loud MESSAGE + return when no
+// Tenstorrent device, when VT_KOLIBRI1_TT_B2II_MODEL is unset, or when the
+// named directory is not a Kolibri-1 checkpoint. Operator-run under the GPU
+// lease; evidence lands in docs/bench-evidence/.
+//
+// Sampling: the per-step minimum the greedy loop needs is the TOKEN ID, not
+// the logits row — the argmax runs on-device (vt::GreedyArgmax, the
+// on-device sampling seam) and only the 8-byte id crosses back. Set
+// VT_KOLIBRI1_TT_DECODE_BENCH_LOGITS_COPY=1 to also copy the full f32
+// logits row per step (diagnostic only; it is NOT the default measurement).
+#include <doctest/doctest.h>
+
+#include <algorithm>
+#include <chrono>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <filesystem>
+#include <fstream>
+#include <memory>
+#include <numeric>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "vllm/model_executor/model_loader/safetensors_reader.h"
+#include "vllm/model_executor/models/kolibri1_shared.h"
+#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/models/model_registry.h"
+#include "vllm/transformers_utils/hf_config.h"
+#include "vt/backend.h"
+#include "vt/device.h"
+#include "vt/dtype.h"
+#include "vt/ops.h"
+#include "vt/tensor.h"
+
+using namespace vllm;  // NOLINT
+
+namespace {
+
+double NowSec() {
+  return std::chrono::duration<double>(
+             std::chrono::steady_clock::now().time_since_epoch())
+      .count();
+}
+
+v1::CommonAttentionMetadata OneReqMeta(int64_t t, int64_t ctx_before,
+                                       int64_t block_table_num_cols) {
+  v1::CommonAttentionMetadata m;
+  m.num_reqs = 1;
+  m.num_actual_tokens = static_cast<int>(t);
+  m.max_query_len = static_cast<int>(t);
+  m.max_seq_len = static_cast<int>(ctx_before + t);
+  m.query_start_loc = {0, static_cast<int32_t>(t)};
+  m.query_start_loc_cpu = m.query_start_loc;
+  m.seq_lens = {static_cast<int32_t>(ctx_before + t)};
+  m.seq_lens_cpu = m.seq_lens;
+  m.num_computed_tokens_cpu = {static_cast<int32_t>(ctx_before)};
+  m.block_table_tensor.assign(static_cast<size_t>(block_table_num_cols), 0);
+  for (int64_t i = 0; i < block_table_num_cols; ++i)
+    m.block_table_tensor[static_cast<size_t>(i)] = static_cast<int32_t>(i);
+  m.block_table_num_cols = static_cast<int>(block_table_num_cols);
+  for (int64_t i = 0; i < t; ++i) m.slot_mapping.push_back(ctx_before + i);
+  m.causal = true;
+  return m;
+}
+
+// Prefill shape: the engine CANNOT prefill a whole 128-token prompt in one
+// forward — one whole-prompt routed dispatch streams every selected expert's
+// slot bytes in a single charge and trips the B1 per-step stream bound
+// (observed here: 1404123840 B against the 1179936000 B bound at layer 3;
+// the b2ii gate records the same refusal, 1242865920 B at layer 15). The
+// one-token incremental walk over the persistent KV cache is the faithful
+// same-math alternative (test_kolibri1_tt_b2ii.cpp device gate, which cites
+// test_kolibri1_w3.cpp:298-310), so the prefill walks the prompt one token
+// per forward and the walk is what the prefill timer covers.
+constexpr int kPromptLen = 128;
+constexpr int kDecodeSteps = 64;
+
+}  // namespace
+
+TEST_CASE("kolibri1 TT decode bench: 128-token prefill, 64 greedy steps, "
+          "wall/prefill/decode split") {
+  if (vt::TryGetBackend(vt::DeviceType::kTENSTORRENT) == nullptr) {
+    MESSAGE("SKIPPED: no Tenstorrent device on this box");
+    return;
+  }
+  const char* model_dir = std::getenv("VT_KOLIBRI1_TT_B2II_MODEL");
+  if (model_dir == nullptr || *model_dir == '\0') {
+    MESSAGE("SKIPPED: VT_KOLIBRI1_TT_B2II_MODEL not set (no real checkpoint "
+            "requested)");
+    return;
+  }
+  const std::string dir = model_dir;
+  if (!std::filesystem::exists(dir + "/model.safetensors.index.json")) {
+    MESSAGE("SKIPPED: " << dir << " is not a Kolibri-1 checkpoint");
+    return;
+  }
+
+  // ---- 1. Load the REAL checkpoint through the production registry path.
+  const HfConfig config = vllm::LoadHfConfig(dir + "/config.json");
+  const ModelRegistration& reg = ModelRegistry::Resolve(config);
+  REQUIRE(reg.architecture == "Kolibri1ForCausalLM");
+  reg.factory->parse_config(config);
+  const auto index = nlohmann::json::parse(
+      std::ifstream(dir + "/model.safetensors.index.json"));
+  std::set<std::string> shard_names;
+  for (const auto& [name, shard] : index.at("weight_map").items()) {
+    (void)name;
+    shard_names.insert(shard.get<std::string>());
+  }
+  std::vector<SafetensorsFile> shards;
+  for (const std::string& shard : shard_names)
+    shards.push_back(SafetensorsFile::Open(dir + "/" + shard));
+  const ModelSource source = ModelSource::FromSafetensors(shards);
+  const double t_load0 = NowSec();
+  std::unique_ptr<LoadedModel> model = ModelRegistry::Load(config, source);
+  REQUIRE(model != nullptr);
+  const double load_s = NowSec() - t_load0;
+  MESSAGE("load " << load_s << " s");
+
+  // ---- 2. Prepare on the TT queue (the production residency the decode
+  // runs) and build the device KV topology.
+  vt::Backend& be = vt::GetBackend(vt::DeviceType::kTENSTORRENT);
+  vt::Queue q = be.CreateQueue();
+  const double t_prep0 = NowSec();
+  ModelRegistry::Prepare(*model, config, q);
+  const Kolibri1TTResidentDeviceContext& ctx =
+      Kolibri1LoadedModelTTContext(*model, q);
+  REQUIRE(ctx.built);
+  MESSAGE("prepare " << NowSec() - t_prep0 << " s, "
+                     << ctx.projections << " projections, "
+                     << double(ctx.uploaded_bytes) / (1 << 30) << " GiB bf16");
+
+  const Kolibri1Weights& w = Kolibri1LoadedModelWeights(*model);
+  const Kolibri1Params& p = w.params;
+  const int64_t hkv = p.num_key_value_heads;
+  const int64_t hd = p.head_dim;
+  const int64_t block_size = 16;
+  // 128-in + 64-out = 192 tokens -> 12 blocks; carry 32 for headroom.
+  const int64_t num_blocks = 32;
+  const int64_t block_table_cols = num_blocks / 2;
+  const int64_t kv_bytes = num_blocks * 2 * block_size * hkv * hd *
+                           static_cast<int64_t>(vt::SizeOf(vt::DType::kBF16));
+  std::vector<std::shared_ptr<void>> kv_keep;
+  std::vector<PagedKvCache> kv;
+  kv.reserve(static_cast<size_t>(p.num_hidden_layers));
+  for (int64_t l = 0; l < p.num_hidden_layers; ++l) {
+    void* buf = be.Alloc(kv_bytes);
+    kv_keep.emplace_back(buf, [&be](void* ptr) { be.Free(ptr); });
+    std::memset(buf, 0, static_cast<size_t>(kv_bytes));
+    PagedKvCache c;
+    c.data = buf;
+    c.dtype = vt::DType::kBF16;
+    c.num_blocks = num_blocks;
+    c.block_size = block_size;
+    c.num_kv_heads = hkv;
+    c.head_size = hd;
+    kv.push_back(c);
+  }
+
+  v1::GDNAttentionMetadata gdn_meta;
+  std::vector<GdnStateCache> gdn_state;
+  // Deterministic prompt, same recipe as the CPU bench: (i*7919+13)%128000.
+  std::vector<int32_t> seq;
+  seq.reserve(kPromptLen + kDecodeSteps);
+  for (int i = 0; i < kPromptLen; ++i)
+    seq.push_back(static_cast<int32_t>((i * 7919 + 13) % 128000));
+
+  // Device-side greedy sampler scratch: the ids tensor lives on device, only
+  // the 8-byte token id crosses back per step.
+  int64_t* ids_host = static_cast<int64_t*>(be.AllocPinned(sizeof(int64_t)));
+  void* ids_dev = be.Alloc(sizeof(int64_t));
+  std::shared_ptr<void> ids_keep(ids_host, [&be](void* ptr) {
+    be.FreePinned(static_cast<int64_t*>(ptr));
+  });
+  std::shared_ptr<void> ids_dev_keep(ids_dev, [&be](void* ptr) {
+    be.Free(ptr);
+  });
+  const bool copy_logits =
+      std::getenv("VT_KOLIBRI1_TT_DECODE_BENCH_LOGITS_COPY") != nullptr;
+
+  auto GreedyNext = [&](const ForwardLogits& fl) -> int32_t {
+    vt::Tensor ids =
+        vt::Tensor::Contiguous(ids_dev, vt::DType::kI64, q.device, {1});
+    vt::GreedyArgmax(q, ids, fl.device_tensor);
+    be.Copy(q, ids_host, ids_dev, sizeof(int64_t));
+    be.Synchronize(q);
+    return static_cast<int32_t>(ids_host[0]);
+  };
+
+  auto Forward = [&](const std::vector<int32_t>& tokens, int64_t ctx_before,
+                     bool pure_decode) -> ForwardLogits {
+    const int64_t t = static_cast<int64_t>(tokens.size());
+    // The b2ii walk convention: a one-token forward IS a decode step,
+    // including the prompt walk (same math over the persistent KV cache).
+    pure_decode = (t == 1);
+    const v1::CommonAttentionMetadata meta =
+        OneReqMeta(t, ctx_before, block_table_cols);
+    std::vector<int32_t> positions(static_cast<size_t>(t));
+    std::iota(positions.begin(), positions.end(),
+              static_cast<int32_t>(ctx_before));
+    ModelForwardInput in{tokens,
+                         positions,
+                         meta,
+                         gdn_meta,
+                         kv,
+                         gdn_state,
+                         config,
+                         q,
+                         /*logits_indices=*/{},
+                         /*num_reqs=*/1};
+    in.pure_decode = pure_decode;
+    in.uniform_query_len = pure_decode ? 1 : 0;
+    return ModelRegistry::Forward(*model, in);
+  };
+
+  // ---- 3. PREFILL: walk the 128 prompt tokens one forward per step over
+  // the persistent KV cache (see the prefill-shape note above), timed; the
+  // last walk step yields the first greedy token.
+  const double t_pre0 = NowSec();
+  ForwardLogits fl;
+  for (int i = 0; i < kPromptLen; ++i) {
+    fl = Forward({seq[i]}, /*ctx_before=*/i, /*pure_decode=*/false);
+    REQUIRE(fl.rows == 1);
+  }
+  int32_t next = GreedyNext(fl);
+  seq.push_back(next);
+  const double prefill_s = NowSec() - t_pre0;
+  MESSAGE("prefill walk " << kPromptLen << " tokens in " << prefill_s
+                          << " s; first token " << next);
+
+  // ---- 4. DECODE: 64 free-running greedy steps, timed from before the
+  // first step to after the last sync.
+  std::vector<int32_t> chain{next};
+  const double t_dec0 = NowSec();
+  for (int step = 0; step < kDecodeSteps; ++step) {
+    const int64_t ctx_before = static_cast<int64_t>(seq.size()) - 1;
+    fl = Forward({seq.back()}, ctx_before, /*pure_decode=*/true);
+    REQUIRE(fl.rows == 1);
+    if (copy_logits) {
+      // Diagnostic arm only: bring the full f32 row back and argmax on host.
+      std::vector<float> logits(static_cast<size_t>(fl.vocab));
+      be.Copy(q, logits.data(), fl.device_tensor.data,
+              static_cast<size_t>(fl.vocab) * sizeof(float));
+      be.Synchronize(q);
+      next = static_cast<int32_t>(std::max_element(logits.begin(),
+                                                   logits.end()) -
+                                  logits.begin());
+    } else {
+      next = GreedyNext(fl);
+    }
+    seq.push_back(next);
+    chain.push_back(next);
+  }
+  const double decode_s = NowSec() - t_dec0;
+  const double wall_s = prefill_s + decode_s;
+  REQUIRE(chain.size() == static_cast<size_t>(kDecodeSteps) + 1);
+
+  std::string chain_str;
+  for (int32_t t : chain) chain_str += std::to_string(t) + ",";
+  MESSAGE("CHAIN: " << chain_str);
+  MESSAGE("wall " << wall_s << " s, prefill " << prefill_s << " s, decode "
+                  << decode_s << " s, tok/s "
+                  << (static_cast<double>(kDecodeSteps) / decode_s));
+  MESSAGE("load " << load_s << " s");
+}