Skip to content

Select the sequential engine for attention-mask graphs and add llm-runner --batch-prototype - #363

Merged
stikves merged 2 commits into
apple:mainfrom
stikves:sukru/public-batched-stage3-runner-v2
Oct 22, 2026
Merged

stikves merged 2 commits into
apple:mainfrom
stikves:sukru/public-batched-stage3-runner-v2

Conversation

@stikves

@stikves stikves commented Oct 22, 2026

Copy link
Copy Markdown
Contributor

This PR makes the batched engine from #335 usable from llm-runner.

What this adds

  • Engine selection: a dynamic graph whose main function takes an attn_mask input now auto-resolves to coreai-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 explicit coreai-pipelined override 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 through lockstepBatchedGenerate and 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's max_batch_size is rejected.

Scope and behavior

  • Bundles without an attn_mask input select the same engine as before.
  • --batch-prototype is off by default and adds no public API.

Validation

  • Build and swift format lint --strict are 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 the attn_mask variant resolution and the rejected pipelined override.
  • On a dynamic-batch Qwen3-0.6B bundle exported with --attention-mask batched --dynamic-batch-size 4:
    • the default variant now runs and gives the same greedy output as an explicit coreai-sequential,
    • --batch-prototype 1, 2 and 4 all run; row 0 is identical at every N and matches a normal greedy run, and a rerun reproduces every row,
    • a row count of 5, a count of 0, and a pipelined override are each rejected with a clear message.
  • Greedy output with the option off is byte-identical to main on the default, coreai-sequential and coreai-pipelined engines, for a fresh dynamic export and an older dynamic bundle.
  • A static-shape bundle still selects the static-shape engine and gives the same greedy output as before.

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 stikves self-assigned this Oct 22, 2026
@stikves
stikves marked this pull request as ready for review October 22, 2026 17:30
@stikves
stikves requested a review from a team as a code owner October 22, 2026 17:30
@stikves
stikves merged commit 31484f7 into apple:main Oct 22, 2026
3 checks passed
@stikves
stikves deleted the sukru/public-batched-stage3-runner-v2 branch October 22, 2026 19:29
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.

2 participants