Heuristic:Bigscience workshop Petals Short Inference Pool Merging
| Knowledge Sources | |
|---|---|
| Domains | Optimization, Performance |
| Last Updated | 2026-02-09 13:00 GMT |
Overview
Short inference requests (<=128 tokens, or <=1 token for NF4) are routed to a merged inference pool that processes all blocks in a single call, eliminating inter-block scheduling overhead.
Description
During inference, the server runs each transformer block through task pools. For short sequences, the overhead of dispatching to separate pools per block dominates the actual computation time. Petals merges the inference pools for all hosted blocks into a single `_MergedInferenceStep` that processes the entire chain without interruption. Requests are classified as "short" based on `batch_size * length_increment`: if this is <= 128 tokens (or <= 1 token for NF4 quantization), the merged pool is used.
Usage
Applied automatically during inference. The NF4 threshold is much lower (1 token) because bitsandbytes does not yet ship an efficient NF4 kernel for parallel forward passes, making merged processing less beneficial for multi-token NF4 batches.
The Insight (Rule of Thumb)
- Action: Pool merging is automatic; no user action needed. Be aware that NF4 limits merged processing to single-token batches.
- Value: `MAX_SHORT_INFERENCE_TOKENS = 128` (general), `MAX_NF4_SHORT_INFERENCE_TOKENS = 1` (NF4).
- Trade-off: Merged pools eliminate inter-block overhead for short requests but cannot batch across different inference sessions in the merged path.
Reasoning
Autoregressive generation typically produces 1 token per step per session, well within the 128-token threshold. The merged pool avoids re-dispatching through the task queue for each block, reducing latency. The NF4 exception exists because bitsandbytes NF4 kernels are not optimized for batched parallel forward passes (noted as a TODO in the code).
Code Evidence
From `src/petals/server/block_functions.py:26-28`:
# TODO: Increase the NF4 threshold once bitsandbytes ships efficient NF4 kernel for parallel forward
MAX_SHORT_INFERENCE_TOKENS = 128
MAX_NF4_SHORT_INFERENCE_TOKENS = 1
Merge decision from `src/petals/server/block_functions.py:199-200`:
merge_max_tokens = MAX_NF4_SHORT_INFERENCE_TOKENS if quant_type == QuantType.NF4 else MAX_SHORT_INFERENCE_TOKENS
can_merge_pools = batch_size * length_increment <= merge_max_tokens
Pool merging setup from `src/petals/server/backend.py:201-213`:
def merge_inference_pools_inplace(backends: Dict[ExpertUID, TransformerBackend]):
"""Replace each backend's rpc_inference pools with a combined pool runs multiple blocks in one call"""
merged_pool = PrioritizedTaskPool(
_MergedInferenceStep(backends),
max_batch_size=first_pool.max_batch_size,
device=first_pool.device,
name=f"merged_inference",
)
for backend in backends.values():
backend.inference_pool = merged_pool