kolibri1 TT: generalize the B1 stream envelope to multi-token forwards; the device-side router op

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

@@ -0,0 +1,19 @@
+ID: ISSUE-LOCAL-01M4NS6D2GCD5WXBFN24XGXYQ4
+Title: Generalize the B1 per-token stream bound to multi-token forwards; record the chunked-prefill evidence
+Row: MODEL-TEXT-kolibri-1-tenstorrent
+State: OPEN
+Kind: task
+GitHub: -
+Mirror: PENDING
+Availability: FULL
+Created: 2026-10-11
+Updated: 2026-10-11
+Closed: -
+
+## Problem
+
+The B1 per-step stream guard admits only one-token forwards; a 128-token whole-prompt prefill refuses at 1.40 GB against the 1.18 GB bound. Chunked prefill was measured (prefill 156.3 s -> 72.1 s) but chain equivalence is not adjudicable on this instrument (chains differ run to run on identical code).
+
+## Resolution
+
+-

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

@@ -0,0 +1,19 @@
+ID: ISSUE-LOCAL-01M4NZA9B3DBKS3MD0M4TGRVES
+Title: kolibri1 TT: device-side router op (kSigmoidBiasTopK) landed gate-verified but not default; no measured decode win
+Row: MODEL-TEXT-kolibri-1-tenstorrent
+State: OPEN
+Kind: task
+GitHub: -
+Mirror: PENDING
+Availability: FULL
+Created: 2026-10-11
+Updated: 2026-10-11
+Closed: -
+
+## Problem
+
+vt::SigmoidBiasTopK removes the per-layer full-tail router-logits download, but the composite's extra eager launches per layer cost as much as the drain (decode 952 vs 924 ms/step median, within the run-to-run band; b2ii drift slightly worse: 4 vs 3 HARD flips). Default stays the host path; the op is opt-in via VT_KOLIBRI1_TT_DEVICE_ROUTER=1. The win needs a captured/graphed region to remove the launch overhead.
+
+## Resolution
+
+-

include/vt/ops.h

@@ -105,6 +105,7 @@ enum class OpId : uint8_t {
   kKdaGatedDeltaRule,
   kKdaChunkPrefill,
   kMoeRouterTopK,
+  kSigmoidBiasTopK,
   kMoeCombine,
   kAttention,
   kAttentionCross,
@@ -2641,6 +2642,15 @@ using VHeadPermuteFn =
     void (*)(Queue&, Tensor&, const Tensor&, const VHeadPermuteArgs&);
 using MoeRouterTopKFn = void (*)(Queue&, Tensor&, Tensor&, const Tensor&,
                                  const MoeRouterTopKArgs&, const Tensor*);
+// The Kolibri-1 router (kolibri1.py:126-142): selection score = RAW logit +
+// e_score_correction_bias (NOT sigmoid-of-logit + bias — that is the
+// DeepSeek noaux_tc path), weight = sigmoid of the UNBIASED logit, NO
+// renormalization, ties break to the LOWER expert index. This is exactly
+// SigmoidLogitAddRouting (kolibri1_shared.h) as a registered device op so
+// the TT arm can run it on card and download only [T,K] ids + weights
+// instead of the full [T,E] logits tail.
+using SigmoidBiasTopKFn = void (*)(Queue&, Tensor&, Tensor&, const Tensor&,
+                                   const Tensor&, int);
 // The trailing float is `routed_scale` — the routed_scaling_factor applied to
 // the ROUTED sum before the shared term is added (see vt::MoeCombine).
 using MoeCombineFn =
@@ -4643,6 +4653,12 @@ void MoeRouterTopK(Queue& q, Tensor& weights, Tensor& indices, const Tensor& log
                    const MoeRouterTopKArgs& args,
                    const Tensor* e_score_correction_bias = nullptr);
 
+// The Kolibri-1 router op (see SigmoidBiasTopKFn above). logits [T,E] float,
+// bias [E] f32, out: weights [T,K] f32 (sigmoid of the UNBIASED logit),
+// indices [T,K] i32 (biased raw-logit selection, ties to the lower index).
+void SigmoidBiasTopK(Queue& q, Tensor& weights, Tensor& indices,
+                     const Tensor& logits, const Tensor& bias, int top_k);
+
 // Weighted scatter-combine of the per-expert outputs (moe-semantics.md §4/§6).
 //   out[t,:] = routed_scale * sum_j weights[t,j] * expert_out[t,j,:]  (f32 accum)
 //              + shared[t,:]                            (when shared != nullptr)

scripts/env-doc-allowlist.txt

@@ -124,6 +124,7 @@ VT_KOLIBRI1_PROFILE
 VT_KOLIBRI1_TT_B2I_PROGRESS
 VT_KOLIBRI1_TT_B2BI_PROGRESS
 VT_KOLIBRI1_TT_B2II_STREAM
+VT_KOLIBRI1_TT_DEVICE_ROUTER
 VT_KOLIBRI1_TT_FETCH_PROFILE
 VT_KOLIBRI1_TT_FILL_VERIFY_ONCE
 VT_KOLIBRI1_TT_MEMO_FROM_SHADOW

src/vllm/model_executor/models/kolibri1_tt_forward.cpp

@@ -484,29 +484,55 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p,
   // Router logits, f32 (kolibri1.py:171-177) — the bf16 gate over the
   // shared residency seam, f32 out like the CPU row's LinearBTRaw.
   DBuf dlog = LinearBTRaw(d, dhn, w.router_gate, t, DType::kF32);
-  std::vector<float> logits(static_cast<size_t>(t * e));
-  dlog.Download(d, logits.data());
-  // PROTOTYPE (ISSUE-LOCAL-01M4N8Y7XPDC74SED0WEDN0YFR): the bias is a STATIC
-  // weight, so the per-step D2H copy + blocking Synchronize here was 50 extra
-  // host round-trips per decode step. Cache the host image once per device
-  // buffer (the resident weights keep stable device pointers for the process
-  // lifetime) and serve the step from the cache.
-  const std::vector<float>& bias = [&] {
+  // The routing itself. The device composite (vt::SigmoidBiasTopK) exists and
+  // is gate-verified, but MEASURED SLOWER at this shape and is NOT the
+  // default: VT_KOLIBRI1_TT_DEVICE_ROUTER=1 opts in (the composite's extra
+  // eager op launches per layer cost at least as much host serialization as
+  // the full-tail [T,E] download they replace). The host path stays default
+  // until a captured/graphed region removes the launch overhead.
+  static const bool kDeviceRouter =
+      std::getenv("VT_KOLIBRI1_TT_DEVICE_ROUTER") != nullptr &&
+      std::string(std::getenv("VT_KOLIBRI1_TT_DEVICE_ROUTER")) == "1";
+  std::vector<float> logits;
+  std::vector<float> bias;
+  Kolibri1HostRouting route;
+  if (kDeviceRouter) {
+    DBuf topw(d, DType::kF32, {t, top_k});
+    DBuf topi(d, DType::kI32, {t, top_k});
+    vt::SigmoidBiasTopK(d.q, topw.t(), topi.t(), dlog.t(),
+                        ResidentWeight(d, w.e_score_correction_bias, {e}),
+                        static_cast<int>(top_k));
+    route.ids.resize(static_cast<size_t>(t * top_k));
+    route.weights.resize(static_cast<size_t>(t * top_k));
+    topi.Download(d, route.ids.data());
+    topw.Download(d, route.weights.data());
+  } else {
+    logits.resize(static_cast<size_t>(t * e));
+    dlog.Download(d, logits.data());
+    // The bias is a STATIC weight, so the per-step D2H copy + blocking
+    // Synchronize here was 50 extra host round-trips per decode step
+    // (the #3447 bias host cache: the resident weights keep stable device
+    // pointers for the process lifetime — cache the host image once per
+    // device buffer and serve the step from the cache).
     Tensor bt = ResidentWeight(d, w.e_score_correction_bias, {e});
     std::lock_guard<std::mutex> g(BiasMutex());
     auto it = BiasCache().find(bt.data);
-    if (it != BiasCache().end()) return it->second;
-    vt::Backend& be = vt::GetBackend(d.q.device.type);
-    std::vector<float> fresh(static_cast<size_t>(e));
-    be.Copy(d.q, fresh.data(), bt.data, static_cast<size_t>(e) * sizeof(float));
-    be.Synchronize(d.q);
-    return BiasCache().emplace(bt.data, std::move(fresh)).first->second;
-  }();
+    if (it != BiasCache().end()) {
+      bias = it->second;
+    } else {
+      vt::Backend& be = vt::GetBackend(d.q.device.type);
+      bias.resize(static_cast<size_t>(e));
+      be.Copy(d.q, bias.data(), bt.data, static_cast<size_t>(e) * sizeof(float));
+      be.Synchronize(d.q);
+      BiasCache().emplace(bt.data, bias);
+    }
+    route = SigmoidLogitAddRouting(logits, bias, t, e, top_k);
+  }
 
   // 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)) {
+  if (!kDeviceRouter && StageDumpOn() && l3dbg(layer)) {
     const int64_t hk = w.router_gate.shape[1];
     Tensor bt3 = ResidentWeight(d, w.router_gate, {e, hk});
     std::vector<uint8_t> tb(static_cast<size_t>(e) * hk * 2);
@@ -569,9 +595,9 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p,
       }
     }
   }
-  Kolibri1HostRouting route =
-      SigmoidLogitAddRouting(logits, bias, t, e, top_k);
-  RouteDump(layer, logits, route, t, e, top_k);
+  if (!kDeviceRouter) {
+    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.
@@ -597,7 +623,8 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p,
       PlanKolibri1TTMoEDispatch(st.pool.policy, layer_req);
   const Kolibri1TTLayerDispatch& ld = disp.layers[0];
   const Kolibri1TTSlotFetchList list = BuildKolibri1TTSlotFetchList(
-      st.pool.policy, ld, weights, st.pool, st.per_token_stream_bytes);
+      st.pool.policy, ld, weights, st.pool,
+      st.per_token_stream_bytes * t);
   VT_CHECK(st.guard != nullptr,
            "kolibri1-tt B2b-ii: the per-step stream guard is missing");
   st.guard->Charge(layer, list.stream_bytes);
@@ -948,15 +975,18 @@ ForwardLogits ForwardKolibri1TTResidentForward(
     Kolibri1TTStreamingDeviceContext* streaming) {
   VT_CHECK(queue.device.type == vt::DeviceType::kTENSTORRENT, kNonTTRefusal);
   VT_CHECK(ctx.built, "kolibri1-tt forward: the device context is not built");
+  const Kolibri1Params& p = weights.params;
+  Dev d{vt::GetBackend(queue.device.type), queue};
+  const int64_t t = static_cast<int64_t>(token_ids.size());
   // A fresh charge window per step: the B1 per-token stream bound is a
-  // PER-TOKEN bound, and the guard refuses LOUDLY past it.
+  // PER-TOKEN bound, and the guard refuses LOUDLY past it. A multi-token
+  // forward (chunked prefill) draws the SAME per-token envelope t times,
+  // so the window scales with the token count — one token's bound is
+  // unchanged, and a chunk never claims a larger per-token allowance.
   if (streaming != nullptr) {
     streaming->guard = std::make_unique<Kolibri1TTStreamBoundGuard>(
-        streaming->per_token_stream_bytes);
+        streaming->per_token_stream_bytes * t);
   }
-  const Kolibri1Params& p = weights.params;
-  Dev d{vt::GetBackend(queue.device.type), queue};
-  const int64_t t = static_cast<int64_t>(token_ids.size());
   const int64_t h = p.hidden_size;
   const int64_t vocab = p.vocab_size;
 

src/vt/cpu/cpu_ops.cpp

@@ -3303,9 +3303,42 @@ void MoeRouterGroupedTopKKernel(Tensor& weights, Tensor& indices, const Tensor&
 
 // §3 router: softmax (f32, over all E) -> greedy top-k (lowest-index tie-break)
 // -> optional renormalize. weights [T,K] f32, indices [T,K] i32.
+// The Kolibri-1 router as a kernel: EXACTLY SigmoidLogitAddRouting
+// (kolibri1_shared.h, kolibri1.py:126-142) — selection score = RAW logit +
+// e_score_correction_bias, weight = sigmoid of the UNBIASED logit, NO
+// renormalization, ties break to the LOWER expert index (strict `>` scan over
+// ascending index). The CPU reference and the executable spec; the TT kernel
+// mirrors it (the op is the seam that removes the kolibri1 TT arm's
+// per-layer full-tail router-logits download).
+void SigmoidBiasTopKKernel(Queue&, Tensor& weights, Tensor& indices,
+                           const Tensor& logits, const Tensor& bias, int top_k) {
+  const int64_t t = logits.shape[0], e = logits.shape[1];
+  const float* l = logits.Ptr<float>();
+  const float* b = bias.Ptr<float>();
+  for (int64_t i = 0; i < t; ++i) {
+    const float* row = l + i * e;
+    std::vector<char> taken(static_cast<size_t>(e), 0);
+    for (int kk = 0; kk < top_k; ++kk) {
+      int64_t best = -1;
+      float best_score = 0.0f;
+      for (int64_t j = 0; j < e; ++j) {
+        if (taken[static_cast<size_t>(j)]) continue;
+        const float score = row[j] + b[j];
+        if (best < 0 || score > best_score) {
+          best = j;
+          best_score = score;
+        }
+      }
+      taken[static_cast<size_t>(best)] = 1;
+      weights.Ptr<float>()[i * top_k + kk] =
+          1.0f / (1.0f + std::exp(-row[best]));
+      indices.Ptr<int32_t>()[i * top_k + kk] = static_cast<int32_t>(best);
+    }
+  }
+}
+
 void MoeRouterTopKKernel(Queue&, Tensor& weights, Tensor& indices, const Tensor& logits,
-                         const MoeRouterTopKArgs& args, const Tensor* bias) {
-  if (args.num_expert_group > 0) {  // W3 grouped-topk (`noaux_tc`) path
+                         const MoeRouterTopKArgs& args, const Tensor* bias) {  if (args.num_expert_group > 0) {  // W3 grouped-topk (`noaux_tc`) path
     MoeRouterGroupedTopKKernel(weights, indices, logits, args, bias);
     return;
   }
@@ -4525,6 +4558,8 @@ struct Registrar {
                    static_cast<VHeadPermuteFn>(&VHeadPermuteKernel)));
     RegisterOp(OpId::kMoeRouterTopK, DeviceType::kCPU,
                reinterpret_cast<void*>(static_cast<MoeRouterTopKFn>(&MoeRouterTopKKernel)));
+    RegisterOp(OpId::kSigmoidBiasTopK, DeviceType::kCPU,
+               reinterpret_cast<void*>(static_cast<SigmoidBiasTopKFn>(&SigmoidBiasTopKKernel)));
     RegisterOp(OpId::kMoeCombine, DeviceType::kCPU,
                reinterpret_cast<void*>(static_cast<MoeCombineFn>(&MoeCombineKernel)));
     RegisterOp(OpId::kAttention, DeviceType::kCPU,

src/vt/op_provider.cpp

@@ -332,6 +332,8 @@ const char* OpNameImpl(OpId op) {
       return "KdaChunkPrefill";
     case OpId::kMoeRouterTopK:
       return "MoeRouterTopK";
+    case OpId::kSigmoidBiasTopK:
+      return "SigmoidBiasTopK";
     case OpId::kMoeCombine:
       return "MoeCombine";
     case OpId::kAttention:

src/vt/ops.cpp

@@ -3451,6 +3451,32 @@ void MoeRouterTopK(Queue& q, Tensor& weights, Tensor& indices, const Tensor& log
       q, weights, indices, logits, args, e_score_correction_bias);
 }
 
+void SigmoidBiasTopK(Queue& q, Tensor& weights, Tensor& indices,
+                     const Tensor& logits, const Tensor& bias, int top_k) {
+  VT_CHECK(logits.rank == 2 && weights.rank == 2 && indices.rank == 2,
+           "sigmoid_bias_topk: logits/weights/indices rank-2");
+  const int64_t t = logits.shape[0], e = logits.shape[1];
+  VT_CHECK(bias.rank == 1 && bias.shape[0] == e,
+           "sigmoid_bias_topk: bias must be [num_experts]");
+  VT_CHECK(IsFloat(logits.dtype) && bias.dtype == DType::kF32,
+           "sigmoid_bias_topk: float logits, f32 bias");
+  VT_CHECK(weights.dtype == DType::kF32 && indices.dtype == DType::kI32,
+           "sigmoid_bias_topk: f32 weights, i32 indices");
+  VT_CHECK(top_k >= 1 && top_k <= e,
+           "sigmoid_bias_topk: top_k must be in [1, num_experts]");
+  VT_CHECK(weights.shape[0] == t && weights.shape[1] == top_k &&
+               indices.shape[0] == t && indices.shape[1] == top_k,
+           "sigmoid_bias_topk: weights/indices must be [T, top_k]");
+  VT_CHECK(logits.IsContiguous() && bias.IsContiguous() &&
+               weights.IsContiguous() && indices.IsContiguous(),
+           "sigmoid_bias_topk: contiguous required");
+  VT_CHECK(logits.device == q.device && bias.device == q.device &&
+               weights.device == q.device && indices.device == q.device,
+           "sigmoid_bias_topk: device mismatch");
+  reinterpret_cast<SigmoidBiasTopKFn>(GetOp(OpId::kSigmoidBiasTopK, q.device.type))(
+      q, weights, indices, logits, bias, top_k);
+}
+
 void MoeCombine(Queue& q, Tensor& out, const Tensor& expert_out, const Tensor& weights,
                 const Tensor* shared, float routed_scale) {
   VT_CHECK(expert_out.rank == 3 && weights.rank == 2 && out.rank == 2,