[jschmied]

Part of #1293 (step 4). On main (#1837 has landed); this PR's change is its two commits, ab9c1ea88 and 34f66f309. The M3 Pro checks are rerun on this head; see the rerun, where three numbers differ from the ones below under memory pressure.

With TF_KV_SHARE=n, Nemotron on Metal keeps prompt states in the core prompt_radix store over the Nemotron KV pool, as CUDA does since #1756. Unset or 0: the host prompt cache as before.

What changes

• efc4cdc07: a Metal lane can adopt a kept prompt's pages. The engines' core re-exports tensorfold.kv_pool, so Metal and CUDA use the same Pool type, and main's cuda_prompt_pages adapter is used as is.
• 875c05684: a kept state holds only Nemotron's Mamba states (Snap.pages). Its attention rows stay in the pool pages a later stream adopts, so a shared prefix is held once.
• Marks on the prompt chunk grid (2,048 rows). The first version kept states at 256-row page-edge marks. Keeping a state at such a mark ended a prefill chunk there, so even cold passes got new chunk seams (990 tokens became 768 + 222), and the gates caught changed tokens (2 of 4 replies against main). Marks now lie on the chunk grid, so a resumed prompt's chunk seams are a cold pass's. The cost: a prompt resumes at a multiple of the chunk size.
• No kernel changes.

Checks (M3 Pro)

• Flag off: equal to main.
• Concurrent equal to solo at 1, 2, 4 and 8 streams (and at 4 without drafts), all equal to main.
• Resumed equal to cold.
• Round robin at a cap of 2 and of 8 states, with shedding and reclaim: byte-equal. At cap 8 the shedding was page-pressure eviction (4 states at once at 8/8, then 1 more with 12 pages free), not only the state cap.
• Idle census: free + kept == total every time.
• Copy against share, round robin: 63k against 115k of 291k prompt tokens cached; kept state about 74 against about 46 MiB, with prefix pages held once.
• Speed, ABBA: TTFT 2k 1.003, 8k 1.000; decode at 1 and 4 streams 1.000.
• lean 0 new; zig build test test-metal-kv-pages test-native: 590/590 steps, 1,312/1,336 tests (24 skipped); contract nemotron: 63 pass gates 1-4 with 1 known finding (main's experts relu2), 29/29 controls red as planted, limits 0 problems.

[ashhart]

@jschmied, could you rebase this onto current main without #1837's commits (#1837 is on main now), and leave it as a draft until you're ready? Thanks.

[jschmied]

Rebased onto main (933b969fa) without #1837's two commits; the branch is now this PR's own two, ab9c1ea88 and 34f66f309 (head). One conflict, in nemotron/kv_pages.zig: main's debug holds assert (from 4675e47d3) stays as it is, and the census takes this PR's comment and its kept-pages check. Adopted pages go into the lease's ids, so holds still covers a resumed stream's rows, and snapshot.zig only calls it when a state copies rows (with Snap.pages it copies none). On Linux (GB10): lean_check 0 new, zig build test, test-kernel-contract test-golden test-server-security test-pull-security fuzz-cortheon and git diff --check pass. The Metal parts don't build on Linux, so the M3 Pro checks in the description still need rerunning on this head. It stays a draft until they do.

[ashhart]

Thanks @jschmied for the clean rebase. We read the two commits. TF_KV_SHARE stays opt-in, there are no kernel changes, and putting the marks on the 2,048-row chunk grid is the right fix for the chunk-seam token changes. With it off, it's main's path. When your Mac checks have been rerun on 34f66f309, please post them: flag off equal to main, concurrent equal to solo, resumed equal to cold, the round-robin cap with shedding, and the ABBA speed. If it helps, we can also run them on an Apple M3/M5 Ultra.

[jschmied]

Thanks @ashhart. Here are the M3 Pro checks rerun on 34f66f309, against main 933b969fa (this PR's parent, with #1837). Every reply is byte-equal and every gate is green. Three numbers differ from the PR body; they are listed first, with the reason for each.

Where this run differs from the body

• Round-robin cached tokens.
  • Share at cap 8: 75,776 of 290,846 prompt tokens, from 17 hits. The body says 114,688.
  • The copy store: 0 hits. This server started under memory pressure (the Mac was in normal use), so the host store kept no states, and the copy-against-share comparison didn't reproduce this time.
  • Cap 2 is identical to the earlier run: 45,056.
• Cap-8 shedding: 4 evictions, all at the 8/8 state cap. No page-pressure eviction happened this time, so the body's "then 1 more with 12 pages free" did not reproduce.
• Decode is noisy: up to 13 % spread inside one arm, because the box was in use.

Gates

• Flag off: equal to main.
• TF_KV_SHARE=4: concurrent equals solo at 1, 2, 4 and 8 streams, and at 4 with --no-drafts. Solo equals main 4/4, and so does turn 2.
• Resumed (4,096 and 18,432 cached tokens): equal to main's cold replies.
• Round robin, every reply equal:
  • cap 2: 25 evictions, idle 292 free + 96 kept = 388;
  • cap 8: 4 evictions, idle 84 free + 304 kept = 388.
• Idle census: 216 lines, free + kept == total on every one, 0 not free.
• lean 424 on both (0 new); fmt and diff-check clean.
• zig build test test-metal-kv-pages test-native: 609/609 steps (main has grown since the body's 590). The first try failed one unrelated GLM test for lack of memory while the 20 GB server ran; it passed on the rerun.
• Contract nemotron: 63 pass gates 1-4, with 1 known finding (main's experts relu2); 29/29 controls red. The one limits problem is the older offline DeepSeek-V4.1 record, outside this diff.

ABBA speed, share against main: medians of 4 servers per arm, 3 repeats each.

cell	share / main
TTFT 2k	1.000
TTFT 8k	1.002
decode, 1 stream	1.017
decode, 4 streams	1.026

• For reference, main's absolute values: TTFT 3.81 s at 2k and 16.43 s at 8k, decode 65.7 tok/s at 1 stream and 60.6 at 4.
• The decode ratios are noise, not a speed-up: one main server ran at 58.7 tok/s against 64.7 to 67.1 for the others.

Since this Mac runs under daily load, a run on your M3 or M5 Ultra would give cleaner speed numbers and the copy-store comparison. We'd gladly take your offer.

[ashhart]

Thanks @jschmied, that's a thorough rerun: every reply byte-equal, every gate green, and the three numbers that moved explained by the box's state, not the code. We'll take the Mac run you suggested: on our M3 or M5 Ultra, on this head, flag off against main, TF_KV_SHARE=4 against the copy store with cached tokens, concurrent and resumed equality, the cap-8 round robin, and an ABBA speed run. We'll post the numbers here. One small thing: main has moved around the kv_prompts import in tensorfold.zig, so it needs a one-line rebase.

[ashhart]

Thanks @jschmied. Here's the Mac run on an M5 Ultra with Nemotron 3.5 Lightning 30B-A3B MLX 4-bit, balanced order, against main:

Arm against main	Prompts, 32,768 tokens	Decode	1 and 4 streams
flag unset	not in this arm	not in this arm	4 streams -0.2% to -0.3%
TF_KV_SHARE=1	+0.0% to +0.1%	-0.4% to -0.0%	-0.5% to -0.0%
TF_KV_SHARE=8	-0.1% to +0.0%	+0.0% to +0.3%	+0.2% to +0.5%

All three read level. On both TF_KV_SHARE arms, 0 of 20 prompt rows differ across runs.

One correction to what we said we'd run: these arms used 1 and 8, not the 4 you suggested, and this table doesn't report the cached-token counts or the cap-8 round robin separately. It answers the question that matters for landing, that the flag off and on costs nothing here.

This is queued to land, and our prep will carry the one-line rebase around kv_prompts.

[ashhart]

Thanks @jschmied. Main has moved under this one. Nemotron's prompt loop now cuts its chunks in prompt_flight.zig (rows()), so the pass.cut line in backend.zig no longer applies. And with main's native/metal.zig at 590 lines, this PR would take it to 607. Could you rebase onto current main, put the store's cut where the chunk end is now decided, and move something out of metal.zig so it stays at 600 lines or fewer? Please paste the Nemotron Metal host tests on the rebased head.

[jschmied]

Thanks @ashhart. All done on 96cec733d, rebased onto main c3f64ca8a, 3 commits:

• The store's cut: prompt_flight.rows() now takes the store's marks (pass.marks, empty without TF_KV_SHARE) and ends a chunk at the first mark past at, after the requested cuts and the decision row and before the cap. It doesn't depend on mark order. Pass.cut and backend.zig's pass.cut line are gone.
• New test in prompt_flight_test.zig, "Nemotron prefill widths: a chunk ends at every TF_KV_SHARE store mark, in any order, beside the other cuts". 29 rows, cuts {3,20}, decision 12, cap 8, marks {24,16,29,40} give widths 3,8,2,3,4,4,5, and every mark within the prompt is a chunk end. With the mark cut disabled it fails.
• metal.zig: 597 lines. It now imports cuda_window.zig's modelContext instead of its own identical copy. xpu_nem.zig still has a third copy, which I left.
• CHANGELOG: on current Unreleased.

Lean checker 210 and 210, new findings 0. zig fmt --check and git diff --check are clean. zig build test --summary all: 715/715 steps, 10,495/10,577 tests passed (82 skipped).

Nemotron Metal host tests on the rebased head (Apple M3 Pro, Zig 0.17.0). The counts are from the first run on the rebased tree; since then only the new test's name changed, and the rerun on 96cec733d is green, from cache:

check-nemotron-prefill-widths-host: Build Summary: 6/6 steps succeeded
test-nemotron-prefill-widths-host: Build Summary: 3/3 steps succeeded; 5/5 tests passed
test-nemotron-round-depth-host: Build Summary: 43/43 steps succeeded; 49/49 tests passed
test-nemotron-mtp-policy-host: Build Summary: 5/5 steps succeeded; 10/10 tests passed
test-nemotron-restored-head-host: Build Summary: 5/5 steps succeeded; 13/13 tests passed
test-nemotron-subset-head-host: Build Summary: 5/5 steps succeeded; 15/15 tests passed
test-nemotron-grants-host: Build Summary: 5/5 steps succeeded; 13/13 tests passed
test-bolt-host: Build Summary: 11/11 steps succeeded; 452/454 tests passed (2 skipped)
test-living-metal-host: Build Summary: 3/3 steps succeeded; 52/53 tests passed (1 skipped)
test-metal-kv-pages: Build Summary: 45/45 steps succeeded; 48/48 tests passed

[ashhart]

Thanks @jschmied. The cut now sits in rows() beside the other cuts, the new widths test is the right check, and metal.zig at 597 is within the limit. Thanks for the host tests on the rebased head. Two small things: the CHANGELOG hunk needs a rebase onto today's Unreleased lines, and could the line say it needs TF_NEMOTRON_METAL_KV_PAGES=1, since TF_KV_SHARE only acts on the paged pool? Optionally, ChunkOffPages in metal.zig could set problem the way the --learn-conversations refusal does.

[jschmied]

Thanks @ashhart. All three are on 2e5cd2cc0, on main 5a45539af:

• CHANGELOG: rebased onto today's Unreleased lines. It's your commit, kept as it was.
• The line: "Nemotron on Metal: with TF_NEMOTRON_METAL_KV_PAGES=1, TF_KV_SHARE=n keeps prompt states in the core prompt_radix store over the KV pool, as CUDA does; TF_KV_SHARE acts only on the paged pool, and unset or 0 keeps the host prompt cache (metal nemotron: TF_KV_SHARE keeps prompt states as core pool pages (stacked on #1837) #1850). Thanks to @jschmied."
• ChunkOffPages: now sets problem before it returns, as the --learn-conversations refusal does, and says why: "TF_KV_SHARE keeps prompts as whole pool pages on the prompt chunks' grid, and this Mac's chunk of {d} rows is not a multiple of the pool's {d}-token pages: serve without TF_KV_SHARE". No test reaches that message, since the chunk can't be set off the page grid from outside (prefillStep is 2048 or 8192).

native/metal.zig is now 600 lines, at the limit.

zig build test: 715/715 steps, 9,553/9,630 tests passed (77 skipped). Lean checker 155 and 155, new findings 0. git diff --check is clean.

[ashhart]

Thanks @jschmied. Here are the two missing cells on the rebased head 96cec73, against its main parent c3f64ca, on an M5 Ultra with Nemotron 3.5 Lightning 30B-A3B MLX 4-bit.

Store	Requests	Prompt tokens	Cached tokens	Fresh-equal replies
copy-main	36	370324	245928	36/36
share4	36	370324	196608	36/36
cap8-round-robin	40	411380	98304	40/40

Every tested reply's token IDs and text equal its fresh run and the other store's reply. Per-request cached tokens are in cached-tokens.tsv. The functional cells used greedy sampling and draft:false.

The server reported 8192-row chunks, so these prompts were longer than that grid.

Cap 8: 12 conversations, 12 hits, 20 evictions; the state count reached 8 and never exceeded it. All 40 replies equal fresh. 112 idle census lines satisfy free + kept == total. The counters do not distinguish page-pressure reclaim from state-cap eviction; we did not force pressure.

With TF_KV_SHARE unset, warm/base/new/new/base, first process excluded:

Cell	New vs main
Prompt 2048	+0.23%
Prompt 8192	+0.04%
Prompt 32768	-0.05%
Decode, 1 stream(s), greedy and seeded	+0.06% to +0.14%
Decode, 4 stream(s), greedy and seeded	+0.02% to +0.17%

All 192 cross-arm prompt/reply keys match. The full report includes both process medians per arm and all speed findings. TF_NEMOTRON_METAL_KV_PAGES=1 was enabled throughout. Share-enabled concurrency/drafting and forced-pressure cases remain unchecked on this head. No source changes or public post were made.

These ran on 96cec733d. Your rebase since keeps that code byte-identical (the same patch apart from CHANGELOG), and the one new commit, 2e5cd2cc0, only words the ChunkOffPages refusal at startup, so the cells stand for the current head. This is all we were waiting on: #1850 is in the landing queue.

[ashhart]

@jschmied, this was in tonight's landing train, but on today's main it takes zig/src/native/metal.zig to 603 lines (main is at 593, and this adds 10 net), past the 600 limit, so the light gate stops it. Could you make room? One option: move the TF_KV_SHARE open, from the chunk % page_tokens refusal to h.host.info_.prompt_cache = true, into a small function beside cuda_prompt_pages.zig's adapter, and inline nemo_disk, which is used once. The rest rebases cleanly, with only the CHANGELOG context. It goes back in the queue as soon as it fits.

[jschmied]

Thanks @ashhart. Done on 3dca2b131, rebased on main cce0bdea7, with only CHANGELOG context in the rebase.

• Room: your option. Adapter.openOnGrid in cuda_prompt_pages.zig now holds the chunk % page_tokens refusal, with the same message and error.ChunkOffPages, then Adapter.open and the grid rules. h.host.info_.prompt_cache = true stays in metal.zig, since the adapter doesn't see the lane host. nemo_disk is inlined. native/metal.zig is 595 lines.
• Behaviour: unchanged by the move. zig build test: 728/728 steps, 9,925/9,998 tests passed (73 skipped). Lean checker 116 and 116, new findings 0. zig fmt --check and git diff --check are clean. The Nemotron host steps and test-metal-kv-pages pass.
• Served on an M3 Pro: Nemotron 3.5 Lightning with TF_NEMOTRON_METAL_KV_PAGES=1 TF_KV_SHARE=4, --prompt-cache-gib 2, greedy, 64 tokens. A 7,459-token request resumed 6,144 tokens from pool pages after a 7,448-token first turn. Its token_sha256 equals a cold server's, and a second resume gives the same.

[ashhart]

@jschmied, an easier way to room: #2425 moves six shared helpers out of zig/src/native/metal.zig, from 593 to 546 lines, and it's in our landing queue. Once it's on main, this fits as is, with one rebase step: your move of modelContext to cuda_window.zig then deletes it from metal_shared.zig instead of metal.zig. If #2425 lands before you get to it, we can do that rebase here.