From fdd0f7fdc0dbcb013e5137d3f5ef0971dc347cbc Mon Sep 17 00:00:00 2001 From: Luca Barbato Date: Sun, 11 Oct 2026 13:15:56 +0200 Subject: [PATCH 1/3] kolibri1 TT: the captured-decode gap analysis and the router-bias prototype MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BACKEND-TENSTORRENT / ISSUE-LOCAL-01M4N8Y7XPDC74SED0WEDN0YFR. The kolibri1 P150 decode step (~1077-1254 ms, 50 layers) never captures: docs/gap-kolibri1-tt-captured-decode.md records the precise gap against the qwen3.5 driver (no decode-graph scaffold, no Warm* hooks, per-layer router Download, host-driven streaming expert fetch with per-miss readback pivot, device RAC/PA arms declined on block_size 16 vs the %32 TILE guard). Verdict: captured decode is NOT feasible for this model without a device-resident MoE redesign (all experts resident, a vt::MoeRouterTopK TT port, grouped expert GEMM) — the geometry decline itself is a config fix (cache block_size 32), not a hard tt-metal constraint. Prototype: the per-layer e_score_correction_bias D2H copy + blocking Synchronize (50 per step, constant bytes) is now cached once per device buffer in kolibri1_tt_forward.cpp. Measured with test_kolibri1_tt_decode_bench (cherry-picked from 4dab09abd, noted as a manual-equivalent copy in the bench header): decode 1152 ms/step vs 1077/1254 baseline across two control runs — inside the +-10% run-to-run noise band, no win; greedy chains are unreproducible run-to-run in every arm (the row's known 7-flip residue). Finding: the bias syncs wait only on the router GEMM tail; the real eager levers are the expert-fetch readback/verify cost, the per-layer full-tail router download, and the block_size-32 + WarmRacIdx un-decline. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai/glm-5.3-flash [maki] --- .../ISSUE-LOCAL-01M4N8Y7XPDC74SED0WEDN0YFR.md | 43 ++++ docs/gap-kolibri1-tt-captured-decode.md | 191 ++++++++++++++++++ .../models/kolibri1_tt_forward.cpp | 31 ++- 3 files changed, 259 insertions(+), 6 deletions(-) create mode 100644 .agents/issues/BACKEND-TENSTORRENT/ISSUE-LOCAL-01M4N8Y7XPDC74SED0WEDN0YFR.md create mode 100644 docs/gap-kolibri1-tt-captured-decode.md diff --git a/.agents/issues/BACKEND-TENSTORRENT/ISSUE-LOCAL-01M4N8Y7XPDC74SED0WEDN0YFR.md b/.agents/issues/BACKEND-TENSTORRENT/ISSUE-LOCAL-01M4N8Y7XPDC74SED0WEDN0YFR.md new file mode 100644 index 0000000000..ace0e1d17c --- /dev/null +++ b/.agents/issues/BACKEND-TENSTORRENT/ISSUE-LOCAL-01M4N8Y7XPDC74SED0WEDN0YFR.md @@ -0,0 +1,43 @@ +ID: ISSUE-LOCAL-01M4N8Y7XPDC74SED0WEDN0YFR +Title: kolibri1 TT: enable captured decode / dispatch batching +Row: BACKEND-TENSTORRENT +State: OPEN +Kind: gap +GitHub: - +Mirror: PENDING +Availability: FULL +Created: 2026-10-11 +Updated: 2026-10-11 +Closed: - + +## Problem + +kolibri1 P150 decode runs eager (~10 device ops/layer x 50 layers, ~960 ms/step, ~1 tok/s) because it lacks the Warm*/capture driver hooks that qwen3.5 implements; both device capture arms decline on kolibri geometry. Produce gap analysis and prototype. + +## Resolution + +- + + +## Evidence (2026-10-11, gap-analysis session, branch row/tt-kolibri-capture) + +Study committed as docs/gap-kolibri1-tt-captured-decode.md. Capture +verdict: NOT feasible for this model without a MoE-arm redesign — the +per-layer router-logits Download (kolibri1_tt_forward.cpp:478) and the +host-driven streaming expert fetch with per-miss readback pivot verify +(:693-701) sit inside any candidate captured region; ttnn trace +prohibits mid-region host reads. The %32 decline +(tenstorrent_paged.cpp:666/:352) is NOT a hard tt-metal constraint: it +is the sdpa_decode TILE validation, and block_size is our cache +allocation — block_size 32 on the TT arm un-declines it (plus a +WarmRacIdx entry for its shapes). + +Prototype (router-bias host cache, removing 50 blocking per-step +syncs): NO measurable win — decode 1152 ms/step vs baseline +1077/1254 ms/step across two control runs (+-10% run-to-run noise; +greedy chains unreproducible run-to-run in all arms, consistent with +the row's 7-flip residue). Finding: the bias syncs wait only on the +tiny router GEMM tail. Re-ranked levers: (1) expert-fetch readback/ +verify cost per slot miss, (2) device routing via a vt::MoeRouterTopK +TT port to remove the per-layer full-tail download, (3) cache +block_size 32 + WarmRacIdx to enable the device RAC/PA arms. diff --git a/docs/gap-kolibri1-tt-captured-decode.md b/docs/gap-kolibri1-tt-captured-decode.md new file mode 100644 index 0000000000..5bb9bdfe3f --- /dev/null +++ b/docs/gap-kolibri1-tt-captured-decode.md @@ -0,0 +1,191 @@ +# Gap analysis: kolibri1 TT captured decode (dispatch batching) + +Row: `BACKEND-TENSTORRENT`. Issue: +`ISSUE-LOCAL-01M4N8Y7XPDC74SED0WEDN0YFR` (in +`.agents/issues/BACKEND-TENSTORRENT/`). Date: 2026-10-11. Branch: +`row/tt-kolibri-capture`. Read-only study plus one contained eager +prototype; no capture wiring landed. + +## Symptom + +kolibri1 P150 decode: ~960 ms/step over 50 layers (~19 ms/layer, ~10 +device ops/layer). The profile is host-serialized, not +silicon-bound. The captured-decode mechanism is default ON since #1625 +and serves other models on this same card, but kolibri1 never captures. + +## How captured decode works (the qwen3.5 reference) + +Capture on TT is a ttnn mesh trace: `BeginCapture` / +`EndCaptureGraph` / `ReplayGraph` map onto `ttnn::begin_trace_capture` +/ `end_trace_capture` / `replay_trace` (`src/vt/tenstorrent/tenstorrent_ops.cpp`, +the `Trace*` family; `SupportsGraphCapture() == true`). TT also +overrides `support_static_graph_mode() == true` +(`src/vllm/platforms/tenstorrent.cpp:98`), so the platform gate is +open. + +The model-side driver is `Qwen3_5DenseDecodeGraph::Step` +(`src/vllm/model_executor/models/qwen3_5.cpp`, ~11900-12900): a +per-padded-size slot state machine `cold (eager) -> warm -> capture -> +replay`: + +- **Driver scaffold**: padded capture sizes (`PadToCaptureSize` / + `BuildPaddedDecode`), one `SizeSlot` per shape holding PERSISTENT + hidden and `[S,vocab]` f32 logits buffers, in-place `Refresh()` of + persistent host inputs, `vt::BreakableGraph` + + `vt::GraphCaptureScope` (kFull) around `ForwardLayers`, replay per + step. +- **Warm\* driver hooks** (all OUTSIDE capture, every step; they + refresh persistent device inputs the captured trace re-reads): + `WarmAttnCosSin`, `WarmPagedKvShadow` (per KV layer), `WarmDecodePos` + (with in-trace `CaptureDecodePosAdvance` plus_one for the R2 + continuation), `WarmRacIdx` (persistent RAC idx tensors + sharded + input; required by the device `TryReshapeAndCacheDeviceDecode` arm), + `WarmPaMeta`, `WarmDecodeIds` / `EmbedDeviceIdsInto` + (`src/vt/tenstorrent/tenstorrent_capture.cpp`, `tenstorrent_device.h:309-448`). +- **Pool discipline**: `MarkStepBoundary` / `StepDemandProfile` (cold + step) / `PreGrowForCapture` / `PinForGraph` (#1380/#2274) so the + captured region does zero allocations. +- **Host-free captured region**: no `to_vector`/`Download`, no host + sync, no host-side data-dependent control flow between ops + (ttnn aborts capture on any device-to-host read inside the region; + the original NO-GO is recorded in + `.agents/specs/tenstorrent-trace-runner.md` ## Outcome). + +Measured reference point: a served captured replay on this P150 class +runs ~12 ms wall / 0.3 ms replay (`.agents/specs/tenstorrent-capture-economics.md`, +#3327), i.e. three orders of magnitude under kolibri1's 960 ms step. + +## What kolibri1 lacks (the hook list, concrete) + +`ForwardKolibri1TTResidentForward` +(`src/vllm/model_executor/models/kolibri1_tt_forward.cpp:903`) is a +hand-rolled eager loop. Missing items, each mapped to its qwen3.5 +counterpart: + +1. **No decode-graph driver**: no size slots, no cold/warm/capture + state machine, no `vt::BreakableGraph` / `GraphCaptureScope`, no + padding to capture sizes, no persistent hidden/logits buffers (every + step allocates fresh `DBuf`s through the whole loop), no + `Replay()`. +2. **No Warm\* hooks at all**: `grep -n 'Warm' kolibri1_tt_forward.cpp` + is empty. The step metadata (`Kolibri1BuildStepInputs`, :973) is + uploaded per step through the generic staging path; nothing primes + persistent device inputs a trace could bake. +3. **Per-layer device-to-host readback in the MoE router** (:478 + `dlog.Download`) plus (before this branch's prototype) a per-layer + bias D2H copy with a blocking `Synchronize` (:483-484). Two blocking + host round-trips per layer, 100 per step. This alone forbids trace + capture: the router readback sits mid-region. +4. **Host-side data-dependent routing + streaming expert fetch**: the + sigmoid-logit-add top-k runs on host over the readback + (`SigmoidLogitAddRouting`, :522); the routed experts then stream from + host through `BuildKolibri1TTSlotFetchList` with host-built jobs, + per-expert H2D staging, a full-slot D2H readback pivot with + byte-exact verification (:693-701), and a CPU dequant per missed + slot. This is host control flow on data the region computes — + uncapturable by construction. +5. **Device RAC / PA decode arms decline on kolibri geometry** + (`src/vt/tenstorrent/tenstorrent_paged.cpp:666` and `:352`): + `if ((d % 32u) != 0u || (bs % 32u) != 0u) return false;` — head_dim + and block_size must be TILE multiples for + `ttnn::paged_scaled_dot_product_attention_decode`. The checkpoint's + cache block_size (16) fails `bs % 32`, and no `WarmRacIdx` entry + exists for its shapes (the RAC arm's precondition check at + `tenstorrent_paged.cpp:766/772` refuses by name). + +### Is the geometry decline hard? + +The `%32` guard is a ttnn `sdpa_decode` kernel validation, hard for a +fixed block_size layout. But block_size is OUR cache allocation, not a +property of the checkpoint: allocating kolibri1's TT KV cache at +block_size 32 satisfies the guard without touching weights or numerics. +head_dim (128 on the real checkpoint) already passes. So: **not a hard +constraint — a padded-layout/configuration fix** (cache block_size 32 +on the TT arm), though the device RAC arm additionally needs a warmed +`WarmRacIdx` entry and TILE-shadow k/v from the device rope path before +it serves. + +## Why full captured decode is NOT feasible here (the honest verdict) + +Items 1, 2, 5 are the ordinary porting work (the driver scaffold + +hooks + block_size fix; the qwen3.5 template exists and is proven on +this card). Items 3 and 4 are architectural to the B2b-ii streaming-MoE +design and cannot be captured around: + +- The routed expert set is data-dependent per token per layer and + arrives through host-built fetch jobs with host-verified readbacks. + A captured trace bakes fixed device pointers and fixed op sequences; + expert streaming changes both every step. +- Making it capturable means the qwen3.5 fused-MoE shape: ALL experts + resident on device, device-side routing (`vt::MoeRouterTopK` — which + has NO TT implementation today; `src/vt/` shows cpu/cuda/rocm/vulkan + only), and grouped expert GEMM on device. kolibri1 has 384 experts x + 50 layers; resident fp8 experts plus dequant working set against the + P150 DRAM budget that the B2b-ii record already describes as + allocated to within ~0.6 MB/bank is the open sizing question, and the + `MoeRouterTopK` TT port plus a grouped-GEMM TT arm are each + multi-day rows on their own. + +Estimate: driver scaffold + Warm hooks + block_size fix alone ≈ 3-5 +device-days, and would deliver ~zero win while items 3/4 still serialize +every layer on the host. Capture only pays after the MoE arm is +device-resident — a campaign, not a task. This satisfies the stop +condition and is the deliverable verdict. + +## What DOES pay first (eager round-trip elimination) + +The step is host-serialized: every blocking D2H in the loop drains the +whole queue. Removing synchronous round-trips needs no capture and +compounds with any later capture work: + +- **Router bias host cache** (PROTOTYPED, this branch): the + `e_score_correction_bias` is a static weight, yet the loop copied it + D2H with a blocking `Synchronize` every layer, every step — 50 + round-trips per step for constant bytes. Now cached once per device + buffer (`kolibri1_tt_forward.cpp`, `BiasCache()`). Measured result + below. +- **Router logits readback** (:478): required by host routing today; + replaced only by the `vt::MoeRouterTopK` TT port (one small D2H of + ids/weights instead, if the streaming fetch stays host-driven). +- **Per-expert readback pivot + CPU dequant** (:693-701): a full + packed-slot D2H plus byte-exact verify per missed slot. Memoized per + slot, so it decays with resident-slot hit rate, but every miss costs + megabytes of D2H. A verify-knob (trust the H2D staging for + non-first-touch) or device-side dequant removes most of it without + touching numerics-critical paths. + +## Prototype result + +`tests/vllm/models/test_kolibri1_tt_decode_bench.cpp` (cherry-picked +clean from 4dab09abd on `row/tt-kolibri-decode-bench`) measured on the +P150, same recipe (tt-smi reset, cache/shm clean, +`TT_METAL_RUNTIME_ROOT=/tmp/umdtrial-install/libexec/tt-metalium`): + +| arm | prefill walk (128 tok) | decode (64 steps) | ms/step | tok/s | +|---|---|---|---|---| +| baseline run 1 (origin/main) | 193.486 s | 68.918 s | **1077** | 0.929 | +| baseline run 2 (origin/main, control) | 149.693 s | 80.225 s | **1254** | 0.798 | +| bias host cache (prototype) | 165.979 s | 73.701 s | **1152** | 0.868 | + +**Verdict: no measurable win.** The prototype sits inside the +run-to-run noise band of the baseline itself (1077-1254 ms/step across +two identical-baseline runs, ~+-10%), and the greedy chains are NOT +reproducible run-to-run in any arm (all three chains differ from the +first token — consistent with the B2b-ii record's state-dependent +7-flip residue; the bench prints, it does not adjudicate). + +**Finding (more valuable than the prototype): the per-layer bias copy +is NOT a decode lever.** Removing 50 blocking `Synchronize` points per +step moved nothing, so those syncs were cheap — each waits only on the +tiny router GEMM tail, not the queue. The step's serialization lives +in the round-trips that wait on the FULL layer tail: the router-logits +`Download` (kolibri1_tt_forward.cpp:478, one full-tail drain per +layer) and the streaming expert path (per-miss full-slot D2H readback + +byte-exact verify + CPU dequant, :693-701). The re-ranked eager levers +are therefore: (1) the expert fetch path's readback/verify cost per +miss, (2) device-side routing (`vt::MoeRouterTopK` TT port) to remove +the per-layer full-tail download, (3) cache block_size 32 to un-decline +the device RAC/PA arms. The prototype edit is kept (it removes real +work and is inert to numerics — cached bytes are identical) but is not +a perf result. + diff --git a/src/vllm/model_executor/models/kolibri1_tt_forward.cpp b/src/vllm/model_executor/models/kolibri1_tt_forward.cpp index 40f929d278..2bdfc6caff 100644 --- a/src/vllm/model_executor/models/kolibri1_tt_forward.cpp +++ b/src/vllm/model_executor/models/kolibri1_tt_forward.cpp @@ -26,8 +26,7 @@ #include #include -#include -#include +#include #include #include #include @@ -462,6 +461,17 @@ DBuf AttentionBlock(Dev d, const Kolibri1AttnWeights& w, const Kolibri1Params& p // streaming context is attached the B2b-i refusal still fires by name — // the refusal remains ONLY for the streaming-disabled arm and genuine // over-capacity/miss-handle failures (which throw by name). + +// PROTOTYPE (ISSUE-LOCAL-01M4N8Y7XPDC74SED0WEDN0YFR): the per-layer +// e_score_correction_bias host image, cached once per device buffer (the +// resident weights keep stable device pointers for the process lifetime). +std::mutex& BiasMutex() { static std::mutex m; return m; } +std::map>& BiasCache() { + static std::map>* c = + new std::map>(); // never destroyed (#1486 pattern) + return *c; +} + DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, const Tensor& dhn, int64_t t, int64_t layer, const Tensor& sh_gate_w, const Tensor& sh_up_w, @@ -476,13 +486,22 @@ DBuf MoeBlock(Dev d, const Kolibri1MoeWeights& w, const Kolibri1Params& p, DBuf dlog = LinearBTRaw(d, dhn, w.router_gate, t, DType::kF32); std::vector logits(static_cast(t * e)); dlog.Download(d, logits.data()); - std::vector bias(static_cast(e)); - { + // 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& bias = [&] { Tensor bt = ResidentWeight(d, w.e_score_correction_bias, {e}); + std::lock_guard g(BiasMutex()); + auto it = BiasCache().find(bt.data); + if (it != BiasCache().end()) return it->second; vt::Backend& be = vt::GetBackend(d.q.device.type); - be.Copy(d.q, bias.data(), bt.data, static_cast(e) * sizeof(float)); + std::vector fresh(static_cast(e)); + be.Copy(d.q, fresh.data(), bt.data, static_cast(e) * sizeof(float)); be.Synchronize(d.q); - } + return BiasCache().emplace(bt.data, std::move(fresh)).first->second; + }(); // The inherited sigmoid-logit-add routing (kolibri1_shared.h, verbatim // from the CPU row): selection on logits + bias, weights = sigmoid of From 09c82e38ce3513582731b26273f80c744a3ff04a Mon Sep 17 00:00:00 2001 From: Luca Barbato Date: Sun, 11 Oct 2026 16:40:57 +0200 Subject: [PATCH 2/3] docs: classify the three eager-lever diagnostic knobs kernel-internal VT_KOLIBRI1_TT_FETCH_PROFILE (the expert-fetch profiler), VT_KOLIBRI1_TT_MEMO_FROM_SHADOW (the memo-from-shadow lever) and VT_KOLIBRI1_TT_FILL_VERIFY_ONCE (the honest no-op experiment) are engineering switches from the eager-lever campaign: allowlist classification, the same one the other TT diagnostic knobs hold. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai/glm-5.3-flash [maki] --- scripts/env-doc-allowlist.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/scripts/env-doc-allowlist.txt b/scripts/env-doc-allowlist.txt index 76c9a13ab7..b1e2e2fe88 100644 --- a/scripts/env-doc-allowlist.txt +++ b/scripts/env-doc-allowlist.txt @@ -124,6 +124,9 @@ VT_KOLIBRI1_PROFILE VT_KOLIBRI1_TT_B2I_PROGRESS VT_KOLIBRI1_TT_B2BI_PROGRESS VT_KOLIBRI1_TT_B2II_STREAM +VT_KOLIBRI1_TT_FETCH_PROFILE +VT_KOLIBRI1_TT_FILL_VERIFY_ONCE +VT_KOLIBRI1_TT_MEMO_FROM_SHADOW VT_KOLIBRI1_TT_STAGE_DUMP VT_KV_ALLOC_LOG VT_ATTN_SELECT_LOG From eadc5d1b8dbf4654185f87ab69d0fada7469a76e Mon Sep 17 00:00:00 2001 From: Luca Barbato Date: Sun, 11 Oct 2026 18:53:05 +0200 Subject: [PATCH 3/3] docs: list the kolibri1 TT captured-decode gap analysis in the site nav check-site's bijection requires every top-level docs/*.md to have a nav.yaml entry; the gap analysis is a published page. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai/glm-5.3-flash [maki] --- website/data/nav.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/website/data/nav.yaml b/website/data/nav.yaml index 76a8c24f4d..ef49cee218 100644 --- a/website/data/nav.yaml +++ b/website/data/nav.yaml @@ -28,3 +28,5 @@ docs: label: SGLang compatibility - file: ROCM.md label: ROCm backend + - file: gap-kolibri1-tt-captured-decode.md + label: Kolibri1 TT captured decode gap analysis