[jschmied]

Part of #1156 (was ac90d6b8 there). Rebased on main 3f144ccf8; #1355 is closed, and its commits are gone.

An MTP stream's ring and its state now sit in its pool lane's run instead of a fixed per-lane allocation. At --parallel 4 / 32768 the 35B with its MTP layer goes from 64 to 0 MiB fixed a lane, bf16 and fp8. The 27B with DFlash2 stays at 23 MiB. The ring is zeroed on the engine's stream and returned only after the drain, as the maintainer fixes on main expect.

Gates on GB10 (Qwen3.6-35B-A3B MLX 4-bit with --drafter for its MTP layer, Qwen3.8-27B MLX 4-bit with DFlash2; bf16 and fp8 KV, --parallel 8, 16384, TF_QWEN_KV_PAGES=256, 37 greedy outputs per arm): flags unset == main 37/37, paging == main paging 22/22, drafted == plain, concurrent == solo at 1/2/4/8, and no lane pages left at idle, on all four configurations. zig build test -Dnvcc=… exits 0, lean shows 0 new.

[jschmied]

Follow-up check on the borrowed ring and state lifetime:

From reading the code, the lane cannot recycle its run while MTP work is outstanding, on any of the four paths:

• Every MTP launch, copy and readback goes on the engine's one stream (e.stream), the same stream that zeroes a lane's run when the lane opens.
• A successful absorb or hold ends by waiting on that stream. Work can only be left queued when a call fails partway.
• The run goes back in exactly one place: Drafted.dropStream → drafter release (host only, no GPU work) → target release → Pages.close. Pages.close drains e.stream before releaseRun, and the drain either completes or exits the process.
• Cancel mid-draft, failed prefill, failed drafting and normal release all end in that dropStream. A cancel is only acted on between prompt cuts or between rounds, never inside an MTP call. Releasing twice is a no-op (lanes.remove and caches.fetchRemove both find nothing the second time).

The fallback context (ring_of.area() null) owns buffers it got from create(). A borrowed context is destroyed on release and never pooled; only owned contexts go back to the pool. Contexts are keyed by stream and made fresh at each open, so nothing ties one to a lane index.

The check did find a real bug, and it sits next to this code rather than in it. The server fills prompts one step at a time (Drafted.prefillStep), and prefillStep opened the drafter before the target's first step. At that point the stream had no lane yet, so ringOf found no cache, and both MTP and DFlash2 took an owned context for every stream. The pool lane's ring went unused, and every stream past the first allocated an unplanned ring at admission. Whole-prompt prefill already opened the drafter after the target, which is why the PR's byte counts and gates looked right. On the server with TF_QWEN_KV_PAGES_ALL, the original stack opened 40 owned drafter contexts in one gate run (35B + MTP); with the fix it opens 0 (27B + DFlash2: 40 → 0 as well).

Fix, as a new commit on this PR: "lanes: a stepped prompt's drafter opens after the target's first step" (core/lanes/drafted.zig). The drafter now opens after the first step and before that step's rows are taken. A failed step releases only the target, since no drafter is open yet; a failed open releases both. The existing stepped-prefill host test now also asserts that each drafter open sees the target's lane: it failed before the fix (0 of 2) and passes after (2 of 2).

Gates (35B-A3B + MTP and 27B + DFlash2, bf16 KV, --parallel 8 at 16384): outputs match the original stack's with every stream on pages (37/37 each); drafted == plain (37/37); concurrent == solo (15/15); flags off == main (37/37); paging == main paging on solo and resumed requests (22/22). zig build test passes, and lean shows 0 new problems.

[ashhart]

Review is on #1355, covering the whole stack #1355–#1359.

[BobClawblaw]

Tested as part of the stack at f3d545559 on an RTX 4090 under Windows (receipt in #1355): with Qwen3.8-27B MLX 4-bit and DFlash2, flags unset at the tip equals the base 21db2c4c9 on 27/27 replies (solo, "draft": false, three at a time, a 6k-token two-turn chat), bf16 and fp8, and drafted == plain 8/8 in every arm, so the lanes change in 03e776237 keeps the DFlash2 replies here. I did not run the 35B with its MTP layer, so this change's own path (an MTP stream's ring in the pool run) is not exercised on this card.

[jschmied]

Added the host test: cuda_mtp_area_test.zig, commit "cuda qwen3_5: a host test of MTP's area in a lane run, borrowed on open and given back on release".

It drives the real Mtp through its drafter vtable (open, release, reserve) without a device. One stream sits on a pool lane whose run holds the area; the other sits on the reserved window. Capacities 1, 3, 4096 and 32768 rows:

• layout: the borrowed ring starts at the run's area and is 2 * capacity * kvWidth * 2 bytes; the newest row's state starts at the next 256-byte boundary, is hidden * 2 bytes, and ends exactly at base + areaBytes(). Both views are marked borrowed.
• a second open of the same stream is refused (StreamOpen).
• release: the borrowed context is freed (the testing allocator's leak check fails otherwise) and the run is left alone. The reserved window's owned context goes back to the pool.
• reserve(8) with lane runs keeps one context of the drafter's own.

Mutation check: a releaseFn that returns the borrowed context to the pool, as the owned ones go, turns the test red. The pool then holds 2 contexts, and the testing allocator reports the leak.

This branch now sits on #1355's new head (the launch-recording test commit). zig build test passes and lean shows 0 new problems.

[jschmied]

Rebased on main 3f144ccf8 (head 0d9734ebf). #1355 is closed, so its three commits are gone from this stack; nothing else changed except that the stepped-open counter of the drafter test moved from drafted_test.zig into drafted_step_test.zig, its only user, to keep the file at 600 lines. Lean 0 new; zig build test passes apart from the two FSR arena tests that b30d5b9d3 fixes on main.

[ashhart]

Thanks for the host test. It covers the area layout, the borrow and release paths and the reserve count we asked
about, and the mutation check shows it would catch the pooling mistake. The drafter-after-first-step fix is right as
well: it matches the whole-prompt path, and a failed open releases only the target.

Your device gates were run on the stack with #1355, before #1373 changed the paged attention. We'll rerun them on
this head: flags unset against main, paging against main paging, and the --alone --serial checks at 1/2/4/8, on
the 27B with DFlash2 and, where we have it, the 35B with MTP. One small thing: the description still says it is
stacked on #1355 with only the last commit new.

[ashhart]

Landed on main with your authorship (469d16b, f3eb483, 816965d). Our GB10 run of the 27B with DFlash2 on the default path: 0 of 20 prompt rows differ from main and the stream sweep is within ±0.4%, with every stream equal to its run alone. Thanks @jschmied. #1357 follows once the share's lane count is fixed, as noted there.

[ashhart]

This is in TensorFold 1.1.8, released today. Thank you, @jschmied!