Skip to content

lanes: an unstarted resumed turn is ranked by its uncached tail - #2276

Merged
ashhart merged 2 commits into
ashhart:mainfrom
devopsfortunato:lanes/cached-tail-admission
Nov 5, 2026
Merged

ashhart merged 2 commits into
ashhart:mainfrom
devopsfortunato:lanes/cached-tail-admission

Conversation

@devopsfortunato

@devopsfortunato devopsfortunato commented Nov 5, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

A prompt that was admitted with a prompt-cache hit but whose pass has not started yet is ranked by its whole prompt length, because remaining() counted only stream.cached, which stays 0 until the backend restores the kept state during the job's first prompt chunk. So a resumed agent turn with a 200K-token history and a 2K new tail counted as 200K rows of work and waited behind long cold prefills (up to the 8-round aging bound), and the ≤512-row short() rule never applied to it.

The admission-time cache hit is already known: stream.reuse.at is set when the lookup finds the kept state (lane_stream.zig). For a job that has not begun filling and whose restore has not failed, remaining() now counts @max(stream.cached, stream.reuse.at) as done, so the turn is ranked by its uncached tail. If the restore then fails (stream.reuse_failed), the backend prefills from 0, so the whole prompt ranks the job again — a failed restore cannot starve other prompts, and the existing aging bound is untouched. Jobs already filling are unchanged (@max(cached, fill_at) when the backend reports fill progress, cached otherwise). Deliberately out of scope: the short() class boundary still uses prompt_len -| cached, so an unstarted resumed turn keeps taking the long-prompt path (aged steps, budget accounting); that is a separate admission path with its own safety argument for a restore that fails mid-pass, and this PR follows the issue's proposal (the ranking estimate) only.

A side effect worth stating: step()'s before/after rows and srptStep()'s decodeFirst sample stop charging a resumed turn's restored prefix as computed rows on its first chunk. That removes an overcount of TENSORFOLD_PREFILL_BUDGET spend and matches how lane_host_load.zig already sizes the same chunk (start = if (reuse_failed) 0 else @max(cached, hint)).

Checks on this head

  • Head commit: fd3581485 (the maintainer's rebase of this branch onto current main; remaining() was byte-identical on every base along the way, the three Zig hunks never moved, and only the CHANGELOG bullet was re-slotted under ## Unreleased, keeping the displaced one). Head 23430fae on top is an empty commit pushed only to retrigger CI after a cancelled runner; it carries no content.

  • Lean checker, python3 tools/zig/lean_check.py, base vs head (the skill's lean_new.py procedure): 425 problems on the base, 425 on the head, 0 new, measured against every base this branch sat on (36f849362 at 5530 files, 09e2f7fa6 at 5532; the final rebase's base adds no findings — the changed files are untouched by it).

  • Tests run, with the command and pass, fail and skip counts:

    • zig build test-lane-api-host --summary all — the real lane host and engine API against the fake backend, no GPU, includes both changed test harnesses: 497 passed, 0 failed.
    • Mutation receipt: with the base remaining() and this head's tests, the same step is 496 pass, 1 fail — the only failure is the new test ("expected 0, found 1"); with the change, 497/497. So the new test fails before the change and passes after it.
    • zig build test --summary all, full suite, on the pre-rebase tree (base fb8d9fcff, same hunks): 492/494 steps succeeded; 4048 passed, 10 skipped, 1 failed. The single failure is roce_inflight_test "the region is DONTFORK and NOHUGEPAGE before registration": madvise(MADV_NOHUGEPAGE) returns EINVAL on this host's kernel (6.18.50+rpt-rpi-v8, aarch64) for every anonymous mapping — reproduced with a plain python ctypes probe on a fresh mmap, no TensorFold code involved, at every mapping size — and the same step passed in the pre-change baseline run on this host earlier the same day. Inherited environment problem, not caused by this change.
    • zig fmt --check on the three changed Zig files and git diff --check: clean.
  • CI on this head (23430fae, the branch rebased by the maintainer onto main with lanes: a reply that ends soon goes first beside a long prompt, and 27B on Metal does so by default #2224 and friends merged): all 7 checks pass — zig (11m52s), zig-windows (17m42s), zig-macos-dev-tools (12m3s), fmt, release-paths (ubuntu-24.04), benchmark-clients (ubuntu-latest), benchmark-clients (windows-latest) (run). Earlier heads were green too: dffb28b86 passed the same 7 checks (run 42099255802); one intermediate run lost zig-windows to the runner's 1-hour cap after a cache miss with no test failure in the log — infrastructure, and the rerun on this head passes in 17m42s.

  • Not run, and why: no Metal, CUDA, HIP or XPU device and no checkpoint on this host, and the change touches no kernel, weight, checkpoint or arithmetic path — device and model gates do not apply; the Windows and macOS host suites are CI's jobs.

Receipt

The change reaches inference only as scheduling order: it decides which waiting prompt pass gets the next chunk. It touches no token, sampling, drafting, arithmetic, storage or wire path, and each stream's pass is unchanged, so the per-stream exactness contracts (drafted = "draft": false, resumed = fresh, concurrent = solo) are order-independent and are pinned by the unchanged suite.

  • Environment: base fb8d9fcff and rebased base 36f849362, head dffb28b86; Zig 0.17.0 per .zig-version; Linux aarch64 host (4 cores, 7 GiB RAM, kernel 6.18.50+rpt-rpi-v8), no GPU, no checkpoint — the lane tests run against lanes.fake.Fake.
  • Paths and platforms: zig/src/core/lane_host_fill.zig is the shared prefill scheduler under every backend's lane host (Metal, CUDA, HIP, XPU and host fake alike), so the ranking change reaches all of them; I ran the host suites on Linux aarch64 (above), CI compiles and runs the host suites on Windows and macOS.
  • Exactness: n/a — no token path is touched; output is unchanged for every admission order. The scheduling contracts that do apply (no starved admission: aging still bounds every wait at preference_rounds; a failed restore ranks by the whole prompt again) are pinned by the new test and the existing aging test.
  • Decode speed: n/a measured — no GPU or checkpoint here, and decode is untouched.
  • Prompt speed: n/a measured — same reason. The intended effect is the issue's own: a cached turn behind a long cold prefill stops waiting as a whole cold prompt (the issue's prior-art engine measured 183 s → 2.3 s for that case; I make no number claim for TensorFold).
  • Licence of any new weights or drafter: n/a.

Checklist

  • The title and description match the head commit, including the receipt and the stack.
  • New tests fail before the change (496/497 with the base remaining(), only the new test failing) and pass after (497/497); the test has no missing dependency to skip.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths in the diff or its commits. No AI attribution lines.

For a change that reaches inference:

  • Every token still goes through the lane rounds — no path added or removed; chunk sizing, share/debt and the admission budget keep their current rules.
  • Drafted equals plain, resumed equals fresh, concurrent equals solo: order-independent by construction here (per-stream passes unchanged), and the unchanged suite's equality tests pass on this head; not re-measured on a device (none available), as declared above.
  • No precision traded for speed: no bits change anywhere.
  • Decode and prompt processing are no slower than main: the ranking key only changes for jobs whose pass has not started and whose restore has not failed — every other state compares exactly as before, and @max(cached, reuse.at) == cached there; aging still bounds every wait.

@ashhart ashhart left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @devopsfortunato. This is a clean fix, and it does what #2256 describes.

We ran it on an M3 Ultra: Flash Next MLX 4-bit MTP, --parallel 2, prompt cache on. A turn resuming a 42.6K-token conversation, with 42,561 tokens cached, was sent 1 s after a cold prompt of about 39.5K tokens. We used two fresh servers for each side and three repeats on each.

First token main this PR
The resumed turn, beside the cold prompt 28.4-28.8 s 5.2-5.4 s
The cold prompt 29.0-29.4 s 29.7-30.0 s
The resumed turn alone 0.21-0.22 s 0.21-0.22 s

Every reply, from the resumed turns and the cold prompts alike, is the same on both sides.

So the resumed turn no longer waits behind the whole cold prompt, and the cold prompt pays about 0.6 s (2%) for it. That's the trade we want. It goes in the landing queue once CI is green on this head.

A job whose prompt pass has not started is ranked by max(cached,
reuse.at) instead of cached alone: the prompt-cache hit is known at
admission, while cached stays 0 until the backend restores the kept
state in the job's first chunk. A long cached history with a short new
tail no longer waits behind cold prefills as if it were a whole cold
prompt. A rejected restore (reuse_failed) ranks the whole prompt again.
The two test harnesses now define the stream fields the ranking reads.
The short() class boundary is unchanged.
@ashhart
ashhart force-pushed the lanes/cached-tail-admission branch from 580401a to 23430fa Compare November 5, 2026 21:55
@ashhart
ashhart merged commit 882ca0a into ashhart:main Nov 5, 2026
7 checks passed
@ashhart

ashhart commented Nov 5, 2026

Copy link
Copy Markdown
Owner

Thank you @devopsfortunato. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

lanes: cached-tail-aware admission — an unstarted resumed turn is ranked by its whole prompt length and waits behind a long cold prefill

2 participants