Skip to content

nemotron metal: prefill runs a chunk a call, so a new prompt no longer stalls other streams' rounds (#3021) - #3054

Open
ashhart wants to merge 2 commits into
mainfrom
lead/nemotron-metal-prefill-step
Open

ashhart wants to merge 2 commits into
mainfrom
lead/nemotron-metal-prefill-step

Conversation

@ashhart

@ashhart ashhart commented Nov 12, 2026

Copy link
Copy Markdown
Owner

Part of #3021 (#3025 did the same for Gemma 4). Nemotron on Metal ran a new stream's whole prompt in one call on the lane thread, so with shared rounds (#3005) every other stream's reply waited for the whole prefill.

Change (families/nemotron/, Metal backend): the backend now declares prefill_step and cancel_prefill beside reports_fill, so the lane host runs one chunk a call and decode rounds run between chunks.

  • Same arithmetic: prefill_steps.zig's chunk is main's loop body: the cut from prompt_flight.rows (planned starts, the lesson decision row, the row cap), then kv_pages.grow and prompt_flight.submit. The whole pass (whole, still the vtable's prefill) and the stepped one (step) both run it, so the chunks are the same. Nemotron declares no chunk invariance, so the host never caps a step (core/lane_host_fill.zig) and nothing moves the cuts.
  • Shared prompt buffer: each step lands the queue and loads its own prompt's ids before encoding, so a chunk never reads another prompt's ids. The whole pass keeps its one-chunk lookahead.
  • Mamba and MTP state: each chunk takes a fresh pool slot and advances the stream's own cache, as before, and the MTP head absorbs each chunk's rows inside that chunk's encode.
  • Cancel between chunks lands the queue and keeps the state at the last chunk's end.
  • Move: the device side (Chunks.encode, the cache open and restore, the lesson trace) moves from backend.zig to prefill_metal.zig, unchanged apart from qualified names. backend.zig goes from 589 to 500 lines.

Tests (test-nemotron-metal-steps-host, new; recording fakes, no GPU): one call submits one chunk at main's cuts, and the calls' chunks equal the whole pass's; the MTP head absorbs rows 0 to n-2 once each, chunk by chunk; a cancel between chunks keeps the state at the last chunk's end; a decode round runs between two chunks of a long prompt. Each fails on main, which has no prefill_step. check-nemotron-prefill-widths-host now also compiles prefill_step and cancel_prefill; test-nemotron-prefill-widths-host is unchanged. 12 of 12 steps and 9 of 9 tests pass on an arm64 Mac; zig fmt --check is clean; lean finds 345 problems, the same as main.

Before landing, on an M5: replies equal to main's for one stream (greedy, a long prompt that isn't a multiple of the chunk, with the MTP head), for two streams where the second sends a 32K prompt mid-decode, and for a resumed turn; cancel and resend; and stream 1's largest gap between tokens during stream 2's 32K fill (main: the whole fill; this: about one chunk and one round), with no slower time to first token.

  • Synthena (Ash's Agent)

…r stalls other streams' rounds (#3021)

The Metal backend now declares prefill_step and cancel_prefill beside reports_fill. The prompt pass moves out of
backend.zig into prefill_metal.zig (the device side, unchanged) and prefill_steps.zig (whole pass and step, generic,
with recording-fake tests). Both run the same chunk at prompt_flight's cuts, so replies are unchanged; each step lands
the queue and reloads its own prompt ids before encoding, since the prompt buffer is shared.
@ashhart

ashhart commented Nov 12, 2026

Copy link
Copy Markdown
Owner Author

Heads-up for landing order: #3045 adds a token-id bound to Nemotron's Metal prompt pass in nemotron/backend.zig. This PR moves that pass into prefill_metal.zig, so the merge drops the check without any error. Whichever lands second should put it in Ops.open, right after limits.check.

  • Synthena (Ash's Agent)

@ashhart

ashhart commented Nov 13, 2026

Copy link
Copy Markdown
Owner Author

Pushed 081faec on top: the fixes for what the M5 matrix found on the first head.

  • A cancel now publishes the kept end of the prompt (the last landed chunk's end) before its snapshot, so a resumed turn's progress matches what was kept (it used to show the start of the last chunk).
  • The shared prompt-id buffer remembers which stream it holds: consecutive steps of one stream keep their ids and one queued chunk instead of draining and reloading each step; another stream's step still drains first. Release clears it.
  • The lane host credits a cancelled pass's confirmed progress before retiring it, and the end timing starts at the actual cancel callback.

No chunk, grid, kernel or arithmetic change. Host tests: the parent fails the kept-cursor, redundant drain/reload and terminal-credit regressions; this head passes them, the step and flight tests and the shared lane and cache suites.

M5 Ultra rerun of the full matrix against main (240 replies, every one byte-identical): every cell level except plain lone decode at 239.04 tok/s against main's 239.12 (-0.03%); main's own two builds were 239.12 and 239.93, so that gap is inside main's spread, not a change.

  • Synthena (Ash's Agent)

This branch has not been deployed

No deployments
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.

1 participant