Repository navigation
Select the sequential engine for attention-mask graphs and add llm-runner --batch-prototype - #363
Merged
stikves merged 2 commits intoOct 22, 2026
Conversation
Auto-detection picked the pipelined engine for every dynamic graph, and the pipelined engine rejects an attn_mask input with a bare input-count error. A dynamic graph whose main function takes attn_mask now resolves to coreai-sequential, which fills the mask. An explicit coreai-pipelined override on such a graph is rejected. The error for any incompatible override now includes the reason from the compatibility check instead of only "incompatible with model structure".
--batch-prototype N runs the prompt as N identical rows through lockstepBatchedGenerate on the sequential engine and prints each row. Row 0 is greedy and matches a normal greedy run of the same prompt. The other rows sample at --temperature (0.8 if that is 0) with a fixed seed per row, so reruns repeat. The option is a diagnostic and off by default. It is checked before the model loads, and a row count above the bundle's max_batch_size is rejected.
stikves
requested review from
Lewis300,
carinapeng,
kevchengcodes,
srjoglekar246 and
tjia1818
October 22, 2026 17:30
stikves
marked this pull request as ready for review
October 22, 2026 17:30
srjoglekar246
approved these changes
Oct 22, 2026
This was referenced Oct 22, 2026
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.
This PR makes the batched engine from #335 usable from
llm-runner.What this adds
attn_maskinput now auto-resolves tocoreai-sequential, the only engine that fills the mask. Before, auto-detection picked the pipelined engine, which fails on these bundles with an input-count error. An explicitcoreai-pipelinedoverride on such a graph is rejected, and incompatible-override errors now include the reason.llm-runner --batch-prototype N: a diagnostic that runs the prompt as N identical rows throughlockstepBatchedGenerateand prints each row. Row 0 is greedy. The other rows sample at--temperature(0.8 if that is 0) with a fixed seed per row. The row count is checked before the model loads, and a count above the bundle'smax_batch_sizeis rejected.Scope and behavior
attn_maskinput select the same engine as before.--batch-prototypeis off by default and adds no public API.Validation
swift format lint --strictare clean. The full Swift test suite passes, apart from the existing Top-K latency timing test, which can fail when the whole suite runs in parallel and passes on its own. New tests cover theattn_maskvariant resolution and the rejected pipelined override.--attention-mask batched --dynamic-batch-size 4:coreai-sequential,--batch-prototype1, 2 and 4 all run; row 0 is identical at every N and matches a normal greedy run, and a rerun reproduces every row,coreai-sequentialandcoreai-pipelinedengines, for a fresh dynamic export and an older dynamic bundle.