Budget inference tokens across workers, not per worker

[anketpratapsingh]

Makes embedding faster and lighter at the same time by changing how the padded-token budget is allocated. Stacks with the #735 work (PR #21): both modify src/main.rs and both carried the same blob_to_f32 change, so a trial merge conflicts (two adjacent new test modules) and merge order matters. Corrected after review — an earlier version of this description wrongly claimed the two were independent.

The problem

TOKEN_BUDGET (4,096) bounds the largest activation tensor one worker allocates, and each worker runs one batch at a time — so peak sidecar memory scaled with workers x per-worker budget. Holding the per-worker budget fixed made memory grow linearly with -j, which is precisely what made parallelism expensive and pushed the host toward low worker counts.

The change

Budget the total across workers instead: TOTAL_TOKEN_BUDGET / workers, floored at 1,024 and capped at the historical 4,096. Peak memory now stays roughly flat as workers scale.

Throughput improves rather than degrading, which is the counterintuitive part: with BatchLongest padding, a larger budget mostly buys wider batches that pad short texts up to the longest item in the batch, so the extra memory was being spent on padded compute.

Measurements

1,453-node repo, Core Ultra 7 165U, bge-small-en-v1.5 / tract. Runs interleaved with cooldowns and repeated, because a first pass at this showed 2x swings from thermal drift on a laptop part; within-set comparisons below are reproducible to under 1 %.

config	total budget	wall	nodes/sec	peak RSS
2 workers x 4,096 (unchanged)	8,192	60.0 s	24.2	307 MB
8 workers x 4,096 (before)	32,768	48.1 s	30.2	634 MB
8 workers x 1,024 (after)	8,192	35.5 s	40.9	314 MB

So at 8 workers: half the memory and faster. Against the 2-worker default, 1.69x the throughput for the same ~310 MB.

Budget sweep at 8 workers (memory scales with the budget as predicted; 512 starts losing to per-batch overhead, which is why the floor is 1,024):

budget	wall	peak RSS
4,096	60.6 s	618 MB
2,048	51.2 s	418 MB
1,024	46.6 s	308 MB
512	50.2 s	245 MB

No quality impact

Batch composition cannot change the vectors — attention masking makes padding inert.

Correction after review: the original verification here used embed-jsonl, which batches at a fixed const BATCH = 32 and never reads TOKEN_BUDGET. Varying the budget therefore could not change its batching, so that check passed trivially and proved nothing about this change. The property itself is real, but the evidence for it is the reindex path: reindexing the same nodes at budgets 4,096 vs 512 — which do produce different groupings in build_batch_ranges — yields bit-for-bit identical stored vectors (independently verified in review, same md5).

Safety of the defaults

• One and two workers keep exactly 4,096, so configurations where the old default was already right do not regress — at 2 workers a smaller budget measured ~9 % slower, which is why this scales rather than simply lowering the constant.
• Single-submitter backends (GPU/ORT) run one inference loop whatever -j says, so they keep the whole budget.
• TRAVSR_EMBED_TOKEN_BUDGET overrides the derivation for tuning.

Tests

Four new tests pin the rule: division across workers, no regression at 1–2 workers, the floor and the single-submitter exemption (including -j 0 not dividing by zero), and the batcher's budget invariant re-asserted at every worker count. Suite: 62 passed / 0 failed on Windows in release mode; clippy -D warnings and rustfmt clean.

Caveats

All numbers come from one low-power laptop CPU and a small repo, so ratios travel better than absolutes. The memory relationship (workers x budget) is hardware-independent; the exact throughput optimum may shift on a machine with more real cores, which is what the env override is for.

Follow-up (not in this PR)

The host derives worker count from P-cores only, so on this hybrid chip it picks 2 and never reaches the 8-worker configuration that this PR makes cheap. Counting E-cores there would unlock the 1.69x by default; it belongs in the travsr repo with its own benchmarks.

[Abhishek5517]

Review: empirical A/B on real data

Built this branch, ran the suite, and measured the memory/throughput and quality claims against a real workload (this repo's graph.db / embed.db, arctic-embed-m-v1.5, dim 768, tract CPU backend, copied to scratch).

Tests

• 62 passed / 0 failed, including token_budget_tests (division across workers, the 1024 floor, the single-submitter exemption incl. -j 0, and the batcher budget invariant at every worker count).

Memory + throughput (same binary, 1,500 pending nodes, only the budget derivation differs)

I isolated the change on one binary by overriding the budget, so nothing but the derivation varies:

config	-j	per-worker budget	wall	peak RSS
pre-change	8	4096 (TRAVSR_EMBED_TOKEN_BUDGET=4096)	81.6 s	2724 MB
this change	8	1024 (derived)	70.4 s	1512 MB
historical default	2	4096 (derived)	122.2 s	2248 MB

• At 8 workers: 45% less peak memory and faster (1512 MB vs 2724 MB, 70.4 s vs 81.6 s). RSS moving with the budget also confirms the non-accelerated batcher path is the one being exercised.
• The real default shift (2 workers to 8): 1.74x throughput at less memory than the 2-worker default. Matches the 1.69x claim in the description.

Absolute numbers differ from the PR (bigger model and repo, different machine), but every relationship the PR asserts holds.

Quality: bit-identical, verified on the reindex path

Reindexed the same 120 nodes at budget 4096 vs 512 and compared the stored vectors: bit-for-bit identical (same md5). Because these two budgets produce genuinely different batch groupings on the reindex batcher, this is a stronger check than the one in the description (see next point).

Two notes

1. The "bit-identical through embed-jsonl" verification does not actually vary the batcher. embed-jsonl batches at a fixed const BATCH = 32 and never reads TOKEN_BUDGET, so changing the budget cannot change its batching and the check passes trivially regardless. The padding-inert property is real (I confirmed it on the reindex path above), but the description should cite the reindex path, not embed-jsonl.
2. "Independent of the #735 work (PR #21) and touches different code" is not accurate. This branch and #21 both modify src/main.rs and both add the same blob_to_f32 as_chunks change to src/model.rs. A trial merge conflicts in src/main.rs (two adjacent new test modules). Trivial to resolve, but the merge order matters and the claim should be corrected.

I could not reproduce the clippy 1.98 chunks_exact_to_as_chunks lint locally (my toolchain is 1.97), so that motivation is taken on trust from CI. The migration is semantically identical.

Verdict

Change is sound and the win is real. No blocking issues, just the two description corrections above. The scoped-out follow-up (host derives workers from P-cores only, so the 8-worker regime is not reached by default on hybrid chips) is the right call and belongs in the travsr repo.

[Abhishek5517]

Approving. Real-data A/B confirms the win (8 workers: 45% less peak RSS and faster; 1.74x over the 2-worker default), and vectors are bit-identical across budgets on the reindex path. See my detailed comment above. Non-blocking: fix the two description claims (the embed-jsonl bit-identical check does not vary the batcher; this is not independent of #21) and resolve the src/main.rs conflict on merge.

[anketpratapsingh]

Both corrections are right, and the first one is a real error on my part — thank you for catching it.

1. The embed-jsonl verification was vacuous. I checked your claim against the source rather than take it on trust: src/bin/embed_jsonl.rs has const BATCH: usize = 32 and chunks with order.chunks(BATCH) at line 197, with no reference to TOKEN_BUDGET or build_batch_ranges anywhere in the file. So varying the budget could not change its batching, and the bit-identical result I reported would have appeared no matter what this PR did. I stated that as verification when it demonstrated nothing. Your reindex-path check (budgets 4,096 vs 512, matching md5) is the evidence that actually bears on the change, since those budgets do produce different groupings in build_batch_ranges.

2. The independence claim was wrong. Confirmed by trial merge: #21 and this branch collide in src/main.rs (both insert a test module immediately before mod query_memo_tests) and both carried the same blob_to_f32 change.

The description is now corrected on both points, marked as post-review corrections rather than quietly edited.

On merge order: this PR landed first, so I rebased #21 onto the new main, kept both test modules, and the duplicate blob_to_f32 commit dropped out automatically. #21 is now MERGEABLE at three commits — rustfmt, clippy -D warnings and 68 tests green locally.

On the clippy 1.98 lint you could not reproduce on 1.97: that is expected, the lint is new in 1.98 and the runners have moved. main was red on it independently of either PR, which is why both branches ended up carrying the same migration.

One thing your review makes obvious that mine did not cover: there is no automated test pinning padding-invariance on the real batcher path — it needs inference, so it belongs with the #[ignore]d model-backed tests that run against staged fixtures. Happy to add one as a follow-up if you think it earns its keep.