| """Batch planning for the server's queue: which queued rows share one forward. |
| |
| Pure python, no torch, no engine: `plan_batches` takes the row lengths and two callables that describe the engine's |
| bucket grid, and returns the groups to run. `decider.serve.batcher` calls it once per collection. |
| |
| Cost model. A forward of B rows padded to length T costs `overhead + B * T` token-units: `overhead` is the fixed cost of |
| one forward (launch, replay, the slot gather and the device-to-host copy) expressed as the number of padded tokens that |
| take the same time, and `B * T` is the padded token work. Two rows of different lengths are therefore worth merging |
| into the longer row's bucket exactly when the padding they add is cheaper than a second forward's overhead. |
| `DEFAULT_MERGE_OVERHEAD_TOKENS` is measured on decider-2b (docs/SERVING.md, section on the batching policy); the server |
| reads `DECIDER_MERGE_OVERHEAD_TOKENS` over it. |
| |
| The partition is exact, not greedy. Rows are sorted by padded length descending (stable, so rows of the same bucket keep |
| their arrival order) and split into consecutive groups; each group runs at the bucket of its longest member, so a group |
| starting at position k costs `overhead + g * T[k]`. A dynamic program over the sorted sequence takes the cheapest split, |
| subject to `g <= min(max_batch, max_rows(T))`, with ties going to the larger group. The per-bucket grouping the server |
| did before is one of the partitions the program may choose (rows of equal length are adjacent in the sorted order), so |
| the planned cost is never above it, and with `overhead = 0` the plan is exactly that grouping. |
| """ |
|
|
| DEFAULT_MERGE_OVERHEAD_TOKENS = 512 |
|
|
|
|
| def group_cap(T, max_rows, max_batch): |
| """Rows one forward may take at padded length T: the engine's widest captured batch bucket, and the server's cap.""" |
| return max(1, min(int(max_batch), int(max_rows(T)))) |
|
|
|
|
| def batch_cost(size, T, overhead=DEFAULT_MERGE_OVERHEAD_TOKENS): |
| """Cost of one forward of `size` rows padded to `T`, in token-units.""" |
| return overhead + size * T |
|
|
|
|
| def plan_cost(groups, overhead=DEFAULT_MERGE_OVERHEAD_TOKENS): |
| """Cost of a plan: the sum over its forwards.""" |
| return sum(batch_cost(len(idx), T, overhead) for T, idx in groups) |
|
|
|
|
| def per_bucket_groups(lengths, pad_len, max_rows, max_batch): |
| """The grouping of 1.1.0: rows of the same padded length, chunked to the cap. The baseline `plan_batches` improves on.""" |
| return _exact_bucket(list(range(len(lengths))), [pad_len(n) for n in lengths], max_rows, max_batch) |
|
|
|
|
| def _exact_bucket(idx, pads, max_rows, max_batch): |
| buckets = {} |
| for i in idx: |
| buckets.setdefault(pads[i], []).append(i) |
| out = [] |
| for T, members in buckets.items(): |
| cap = group_cap(T, max_rows, max_batch) |
| for k in range(0, len(members), cap): |
| out.append((T, members[k:k + cap])) |
| return out |
|
|
|
|
| def plan_batches(lengths, pad_len, max_rows, max_batch, overhead=DEFAULT_MERGE_OVERHEAD_TOKENS, mergeable=None): |
| """Partition queued rows into forwards. |
| |
| lengths: token count of every queued row, in arrival order. |
| pad_len(n): the engine's padded length for a row of n tokens (`EngineV2.pad_len`). |
| max_rows(T):rows the engine runs in one forward at padded length T (`EngineV2.max_rows`). |
| max_batch: the server's own cap (`DECIDER_MAX_BATCH`). |
| overhead: fixed cost of a forward in token-units (see the module docstring). |
| mergeable(n): False for a row that must not be padded into another row's bucket. Rows above the last captured |
| length bucket run eager at a request-specific shape, so they are grouped by exact length as before. |
| |
| -> [(padded length, [row indices]), ...]. Every row appears in exactly one group; a group's padded length is at |
| least every member's own padded length; a group holds at most `min(max_batch, max_rows(T))` rows. |
| """ |
| n = len(lengths) |
| if n == 0: |
| return [] |
| pads = [pad_len(x) for x in lengths] |
| merge = [True] * n if mergeable is None else [bool(mergeable(x)) for x in lengths] |
| out = _exact_bucket([i for i in range(n) if not merge[i]], pads, max_rows, max_batch) |
|
|
| order = sorted((i for i in range(n) if merge[i]), key=lambda i: -pads[i]) |
| m = len(order) |
| dp = [0] * (m + 1) |
| take = [0] * (m + 1) |
| for k in range(m - 1, -1, -1): |
| T = pads[order[k]] |
| cap = min(group_cap(T, max_rows, max_batch), m - k) |
| best, best_g = None, 1 |
| for g in range(1, cap + 1): |
| c = overhead + g * T + dp[k + g] |
| if best is None or c <= best: |
| best, best_g = c, g |
| dp[k], take[k] = best, best_g |
| k = 0 |
| while k < m: |
| g = take[k] |
| out.append((pads[order[k]], sorted(order[k:k + g]))) |
| k += g |
| return out |
|
|
