Repository navigation
kolibri1 TT: generalize the B1 stream envelope to multi-token forwards; the device-side router op - #3448
Merged
Conversation
…en forwards The B1 per-step stream guard was a one-token window: the forward charged every layer's fetch against per_token_stream_bytes, so any multi-token prefill forward refused loudly (whole-prompt 128: 1,404,123,840 B against the 1,179,936,000 B bound at layer 3; the b2ii gate recorded 1,242,865,920 B at layer 15). The bound is a per-token envelope, not a per-step buffer: the guard now scales the window by the forward's token count, both the per-forward charge window (kolibri1_tt_forward.cpp) and the per-layer check inside BuildKolibri1TTSlotFetchList. One-token forwards are unchanged — the scale factor is 1. Chunked prefill was measured with the generalization (8-token chunks, bench arm): the 128-token prefill walk dropped 156.3 s -> 72.1 s with the decode split unchanged. It is NOT adopted as the default: the greedy chain could not be proven equivalent, because this instrument is not bit-reproducible run to run — three runs of the same one-token walk gave three different chains (first token 346, 30686, and 10 under chunking), and the b2ii golden gate's HARD flips move between runs (it is a band gate, not a bit gate). Adopting chunked prefill owes band-class equivalence evidence of its own; the decode bench comment records the evidence and the decode bench keeps the faithful one-token walk. Gate evidence (P150, tt-smi reset before every run): - decode bench, chunked arm: prefill 72.1 s vs 156.3 s baseline; decode 59.1 s vs 59.8 s (within run-to-run noise). - b2ii device gate: 15/16 both on origin/main and with this change (the failing golden-chain case flips different prompts per run on origin/main alone) — verdict class unchanged. - test_tenstorrent_backend: 104/106, failing exactly the two known kGdnDecode and kMatmulBTQuantGrouped cases, unchanged. Owes: the chunked-prefill equivalence adjudication (ISSUE-LOCAL-01M4NS6D2GCD5WXBFN24XGXYQ4). FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai/glm-5.3-flash [maki]
lu-zero
force-pushed
the
row/tt-kolibri-b1-routing
branch
from
October 11, 2026 17:46
bd195ae to
f71b7e8
Compare
…gate-verified, measured, not default New shared op vt::SigmoidBiasTopK (OpId::kSigmoidBiasTopK): the Kolibri-1 router of kolibri1.py:126-142 as a registered kernel — selection score = RAW logit + e_score_correction_bias (NOT the DeepSeek noaux_tc sigmoid+bias form; vt::MoeRouterTopK cannot express it, its contract routes bias through the grouped path), weight = sigmoid of the UNBIASED logit, no renormalization, ties to the lower expert index. The CPU kernel (cpu_ops.cpp) is the exact port of SigmoidLogitAddRouting and the reference; the TT kernel (tenstorrent_ops.cpp) is a device composite — broadcast add, top-k via ttnn::experimental::topk_large_indices, dim-1 gather of the unbiased logit, ttnn::sigmoid — that downloads only the [T,K] ids + weights instead of the full [T,E] logits tail. Substrate constraints recorded in the kernel: topk_large_indices (the ONLY top-k op on this tt-metal) requires ROW_MAJOR BFLOAT16 input and k in multiples of 16, so the selection rounds the biased scores to bf16 (the router GEMM is already bf16 compute; the weights stay f32) and slices a k-bucket down to top_k. RED-FIRST evidence: the focused unit gate (test_tenstorrent_backend.cpp "kTENSTORRENT kSigmoidBiasTopK matches the CPU reference") ran red before the TT registration — REQUIRE(OpRegistered) failed — and green after: expert SETS equal to the CPU reference on random [4,384] logits, each weight within 2e-3 of sigmoid(unbiased logit at its own id) AND of the CPU weight for the same id. Recorded deviation: the ORDER within the top-k can swap on bf16 near-ties (selection precision), benign to the model math (the combine sums weight x expert over the k slots; the dispatch dedups). kolibri1_tt_forward.cpp wires the op into MoeBlock behind VT_KOLIBRI1_TT_DEVICE_ROUTER=1; the host path stays DEFAULT because the composite MEASURED NO WIN at this shape (P150, 2026-10-11, tt-smi reset recipe): decode median 60.9 s/64 steps (952 ms/step) over three runs against the host path's 59.1 s (924 ms/step) in-session — inside the run-to-run band, with the b2ii teacher-forced drift slightly WORSE (4 HARD flips vs the default path's 3-4, gaps -7.1/-7.4/-12.7/-4.4). The extra eager op launches per layer cost what the full-tail drain saved; the lever needs a captured/graphed region to amortize the launches. Gate evidence (P150, reset before every device run): - b2ii device gate: 15/16 BOTH on origin/main and with this change; the failing golden-chain case and its HARD flips move run to run on origin/main alone. Verdict class unchanged (issue: ISSUE-LOCAL-01M4NZA9B3DBKS3MD0M4TGRVES records the not-default state). - test_tenstorrent_backend: the new gate passes; the suite is 104/107 — the two known kGdnDecode and kMatmulBTQuantGrouped failures plus one order-dependent RAC flake that passes in isolation and also appeared in the red run without this change's TT kernel. - test_kolibri1_tt 16/16, test_kolibri1 27/27. - decode bench, host path: decode 57.7-59.8 s/64 steps; device router: 60.1/60.9/67.2 s. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai/glm-5.3-flash [maki]
lu-zero
force-pushed
the
row/tt-kolibri-b1-routing
branch
from
October 11, 2026 17:48
f71b7e8 to
7454f07
Compare
…k CPU arm shifted cpu_ops.cpp The new op's CPU reference arm added ~33 lines ahead of DFlashBlockAttentionKernel; the anchor ratchet correctly flagged the stale line. Re-pointed 3619 -> 3652; check-agent-record green. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai/glm-5.3-flash [maki]
The device-router opt-in from the kSigmoidBiasTopK commit is an engineering A/B switch, same class as the other TT diagnostic knobs on the allowlist. Pre-push checks (env-doc + agent-record) are now run before pushing. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:zai/glm-5.3-flash [maki]
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two results from the kolibri1 P150 dispatch-overhead work.
The B1 per-token stream envelope (
kolibri1_tt.cpp:292: 50 layers x topk 6 x 3,933,120 expert bytes = 1,179,936,000 B/forward — a RAM-bandwidth envelope, not a physical buffer) is now scaled by the forward's token count, identity at t=1. That admits multi-token forwards and makes CHUNKED PREFILL implementable: an 8-token-chunked prefill walk measured 156.3 -> 72.1 s (2.2x). NOT adopted in this change: the greedy chain moves, and the investigation surfaced the decisive fact that the P150 chain is not bit-reproducible run to run even on IDENTICAL code (three same-code runs: first tokens 346, 30686, 10; the b2ii gate's hard flips move between runs on origin/main alone). Chain identity is therefore not an adjudicable criterion on this path — the gate is a band gate, and adopting chunking owes band-class equivalence evidence of its own, recorded in ISSUE-LOCAL-01M4NS6D2GCD5WXBFN24XGXYQ4. Gates: b2ii 15/16 both before and after (the failing case moves run to run on main alone); TT suite 104/106 = the 2 known failures.The device-side router op: new shared
vt::SigmoidBiasTopK(OpId::kSigmoidBiasTopK) — kolibri's exact routing, selection on the RAW logit plus bias, whichvt::MoeRouterTopK's noaux_tc contract cannot express. TT arm implemented, unit-tested (127 new assertions), wired behind the not-default state recorded in ISSUE-LOCAL-01M4NZA9B3DBKS3MD0M4TGRVES. Measured honestly: decode 57.7-59.8 s host vs 60.1-67.2 s device-router — NOT faster on the host path today. Its value is structural: it removes the per-layer full-tail logits download that blocks trace capture, and it is a prerequisite step the capture path needs.Also recorded: the bench's greedy chain is not bit-reproducible run to run in ANY arm (consistent with #3447's record); test_tenstorrent_backend has one order-dependent RAC flake that passes in isolation and appeared in the red run without this change's kernel.
FOLLOWING_AGENTS_PROTOCOL
Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: AGENT:zai/glm-5.3-flash [maki]