Repository navigation
Conversation
…ger stalls other streams' rounds MiniMax-M3's Metal backend gains prefill_step and cancel_prefill, and now reports its fill (#3021). The prompt pass moves out of backend.zig into prompt_pass.zig; warm() keeps its stream-less feed. Each call runs one chunk at exactly the whole pass's cuts (the prompt chunk cap, then a planned start or kept mark inside the chunk); MiniMax declares no chunk invariance, so max_rows is ignored and the cuts never move. The whole pass and the steps share one chunk body, so replies are unchanged. Each step drains other streams' steps in flight before its forward, as the whole pass drained before it began. For a drafting stream that steps, each chunk's tapped rows are served by features() while no forward has rewritten that scratch, and a mark reached by the chunk is kept at the next call, after the drafter has absorbed it. A cancel keeps the state at the last chunk's end; a stream that leaves mid-pass is forgotten; image rows are made once, at the first call. Host tests on recording fakes: stepped chunks equal the whole loop's, a resumed pass cuts as main does, a cancel keeps the last chunk's end, another stream decodes between two chunks, and a drafter absorbs every chunk with the mark kept after its absorb. The new step joins the macOS host CI list.
This branch has not been deployed
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
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.
What this changes
MiniMax-M3 on Metal ran a new stream's whole prompt in one call on the lane thread, so with shared rounds every other stream's reply waited for the whole prefill. The backend now declares
prefill_stepandcancel_prefillbesidereports_fill, so the lane host runs one chunk a call and decode rounds run between chunks, as Gemma 4 (#3025), Nemotron (#3054) and LFM2-MoE (#3058) do.Same cuts: one chunk a call at exactly today's cuts. MiniMax-M3 declares no chunk invariance, so the host never caps a step and the cuts never move. The whole pass and the stepped one share the same chunk body.
Each step first waits for other streams' steps still on the GPU. A stepping drafter (DFlash2 or DSpark) gets each chunk's rows through
features(), valid until the next forward, and a mark is kept after the drafter has absorbed its chunk, as lfm2_moe metal: prefill runs a chunk a call, so a new prompt no longer stalls other streams' rounds (#3021) #3058 does for LFM2-MoE.Cancel between chunks keeps the state at the last chunk's end; a stream that leaves mid-pass is forgotten.
Sizes:
backend.zig529 → 467 lines with the pass moved to a newprompt_pass.zig(181); newprefill_steps.zig(52) and its test (292). The new host step is in the macOS host CI list.On cost: an M5 matrix on gemma4 metal: prefill runs a chunk a call, so a new prompt no longer freezes other streams' decode (#3021) #3025 showed lone decode level with main (earlier gaps were run-to-run noise), and a new prompt's first token arriving later while other streams decode is the interleaved rounds, the point of Gemma 4 on Metal: a new prompt freezes every other stream's decode for its whole prefill (no prefill stepping) #3021.
Issue: Refs Gemma 4 on Metal: a new prompt freezes every other stream's decode for its whole prefill (no prefill stepping) #3021.
Stacked on: main. Small adjacent-line conflicts with lfm2_moe metal: prefill runs a chunk a call, so a new prompt no longer stalls other streams' rounds (#3021) #3058 in
prefill_widths_metal.zig/host_ci.zig; whichever lands second rebases.Checks on this head
zig fmt --check: clean.check-metal-prefill-rest-hostpasses (it compiles the real prefill, prefill_step and cancel_prefill bodies).Receipt
Checklist