fix(embeddings): batch internally so a long-chunk request cannot ask for 200GB
Lifting the hardcoded 512 cap (vmlx#255) made a batch of long inputs a single padded forward: 100 chunks of 40k tokens became ~3.3M padded tokens and asked Metal for 200GB, failing the request. That failure mode did not exist before the fix, so the fix was incomplete. The endpoint now splits a request into sub-batches with a bounded padded-token budget (256k, VMLX_EMBED_MAX_PADDED_TOKENS) and clears the allocator between them. Same vectors, same order, no refusal, bounded memory. Group planning estimates length from characters rather than tokenizing twice, and over-estimates on purpose so groups stay small. Measured on Qwen3-Embedding-0.6B: the 100x40k case now completes with all 100 vectors at 1.4GB RSS (was a hard failure), and ordinary work got faster from the per-batch cleanup - 500 chunks 12.8s -> 5.3s, a single 32k-token input 13.6s -> 9.1s. Also corrects an earlier test of mine that counted raw tokenizer calls: batching legitimately changes that count, so it now asserts what it actually claims - exactly one row is re-tokenized when one row overflows.
J
Jinho Jang committed
1bca990966cf6826ee685fcaa579c442090a8986
Parent: 4f2c4ca