AMOOOMA opened a new pull request, #40905:
URL: https://github.com/apache/beam/pull/40905
Split out of #40760, following review feedback to separate the generic
pieces from the vLLM-specific ones. This PR is the generic prefix-aware
batching; vLLM support builds on it in a follow-up.
## Why
`BatchElements` batches per worker, so an LLM model handler behind
RunInference sends its server batches that mix unrelated prompt prefixes, and
the server's prefix (KV) cache gets thrashed instead of reused.
## What
**`BatchElements(key_fn=..., num_shards=...)`** (`transforms/util.py`)
- With `key_fn` set, a new `WithBatchKey` DoFn keys each element *without*
the per-worker UUID that `WithSharedKey` / `WithLengthBucketKey` add.
`GroupIntoBatches` then routes elements that share a key to the same worker, so
every batch holds a single key.
- `num_shards` (an int, a per-key mapping, or a callable) spreads a hot key
over a bounded number of random sub-keys, so one key isn't capped at one worker.
- Incomplete keyed batches flush after `max_batch_duration_secs`, or 1s when
unset (`GroupIntoBatches` treats 0 as "no limit", which would stall low-traffic
keys in streaming).
- Keyed batching works in streaming, unlike the unkeyed in-bundle path.
**RunInference plumbing** (`ml/inference/base.py`)
- `ModelHandler(batch_key_fn=..., batch_key_num_shards=...)` forwards to
`BatchElements`.
**`apache_beam.ml.inference.prefix_batching`** (new, generic for LLM
handlers)
- `PrefixBatchingConfig`, plus default extractors `PromptPrefixKeyFn`
(completions) and `ChatPrefixKeyFn` (OpenAI-style chat messages; unknown
message schemas raise rather than collapsing onto one key).
- Defaults to **8 shards per prefix**. On a 5–40 RPS sine load (OpenOrca,
T4, up to 50 workers), 4 shards left 18–33k requests queued behind hot prefixes
at the end of the load. 8 shards wrote all 160k published requests and still
scaled the pool down to 4 workers at low traffic.
- The generic `BatchElements` / `ModelHandler` defaults stay at 1 shard.
## Testing
- `util_test.py`: homogeneous keyed batches, shard specs (int / mapping /
callable), validation, default buffering.
- `prefix_batching_test.py`: extractors and config.
- `base_test.py`: `batch_key_fn` / `num_shards` reach `BatchElements`, and
an end-to-end RunInference run produces prefix-pure batches.
------------------------
Thank you for your contribution! Follow this checklist to help us
incorporate your contribution quickly and easily:
- [x] Replace these instructions with a description of your change and why
you have made it.
- [x] Mention the appropriate issue in your description (for example:
`addresses #123` or `fixes #456`).
- [x] Update `CHANGES.md` with noteworthy changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
