Conversation
decide_many(contexts) previously ran the full engine path per context: N prefills + N scoring passes (N x rows forward rows). This PR batches the scoring: N prefills + ONE scoring pass over N*R rows. - engine.py: _build_schema_rows(schema, tokenizer, scoring) — the row build (rows/row_field/row_branch/row_option/row_count/tries/row_decision/lead_in) extracted from run_parallel_generation. Rows depend only on schema + tokenizer + scoring, NOT on the context, so every context in a batch shares them. run_parallel_generation now calls the helper (single source). - engine.py: _score_rows gains cache_slots: list | None = None. When given, each chunk merges exactly its own cache slots (row i pairs with slot cache_slots[i]) instead of broadcasting ONE prefill cache. None (default) is the original per-context behaviour, unchanged. - engine.py: run_parallel_generation_batched(model, tokenizer, contexts, ...) — prefills each context (N width-1 passes), builds the shared row set once, runs _score_rows ONCE over all N*R rows with per-row cache slots (BatchKVCache.merge left-pads the different prompt lengths so each slot sees only its own history), splits the per-row logits back per context (re-keyed 0..R-1) and assembles each result through run_parallel_generation(_prebuilt=...) which skips prefill+scoring. NOTE: the per-context caches are passed through UNMERGED — re-merging an already-batched BatchKVCache fails (its offset is an array, not an int); the single merge happens inside _score_rows. - engine.py: run_parallel_generation gains _prebuilt dict (cache, base_ids, t_prefill, scored, t_suffix_eval) — when set, prefill and scoring are skipped and the given cache/scored result is used. - api.py: _assemble_decision extracted from _decide_once (post-processing shared by decide and the batched decide_many); decide_many now calls run_parallel_generation_batched + _assemble_decision per result. - tests/test_w3f_batch.py (6): batched == separate calls (parsed values + per-field probabilities identical to <1e-9), input order preserved, shared pass telemetry, empty contexts [], constraints flow through, decide_many end-to-end with the fake engine. Why batch=1 parity holds: identical rows, identical per-context cache state (left-padding sits inside the causal mask), only the batch width differs. 531 passed, ruff clean.
Owner
Author
|
Superseded by #35 (staged engine, bounded context groups). |
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
1 participant
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.
W3-F brief: decide_many(contexts) previously ran the full engine path per context — N prefills + N scoring passes. This PR batches the scoring: N prefills + ONE scoring pass over N*R rows.
Design
_build_schema_rows(schema, tokenizer, scoring): the row build extracted fromrun_parallel_generation(single source). Rows depend only on schema + tokenizer + scoring — NOT on the context — so every context in a batch shares them._score_rows(..., cache_slots=None): when given, each chunk merges exactly its own cache slots (row i pairs with slot cache_slots[i]) instead of broadcasting one prefill cache. None = original behaviour, unchanged.run_parallel_generation_batched(...): prefill each context (N width-1 passes) → build the shared row set once → ONE_score_rowspass over all N*R rows (BatchKVCache.merge left-pads different prompt lengths, so each slot sees only its own history) → split per-row logits back per context (re-keyed 0..R-1) → assemble viarun_parallel_generation(_prebuilt=...).run_parallel_generation(_prebuilt=...): prefill + scoring phases skipped when the batched caller supplies cache + scored result.api.py:_assemble_decisionextracted from_decide_once; decide_many → batched path.Why batch parity holds: identical rows, identical per-context cache state (left-padding sits inside the causal mask), only the batch width differs. Verified: batched parsed values + per-field probabilities identical to separate calls (<1e-9) on the deterministic fake model.
Note: per-context caches are passed through UNMERGED — re-merging an already-batched BatchKVCache fails (its offset is an array, not an int); the single merge happens inside _score_rows.
Tests (6 in test_w3f_batch.py): batched == separate calls, input order preserved, shared pass telemetry, empty contexts, constraints flow through, decide_many end-to-end (fake engine).
531 passed, ruff clean.