Workspace/Coding labs
Loading progress

Plan cache capacity for variable-length requests

Intermediate55 min

Implement plan_cache(lengths, layers, kv_heads, head_dim, bytes_per_element, budget_bytes), estimating packed per-request KV bytes and selecting the longest arrival-order prefix that fits.

Your task

  1. Require lengths to be a list of nonnegative exact integers; each length already includes reserved output tokens.
  2. Require layers, kv_heads, head_dim, and bytes_per_element to be positive exact integers, and budget_bytes to be a nonnegative exact integer.
  3. Calculate per-token bytes as 2*layers*kv_heads*head_dim*bytes_per_element.
  4. Admit requests in arrival order until the first request that would exceed the budget, then stop even if a later smaller request could fit.
  5. Return per_request_bytes, admitted_indices, used_bytes, and remaining_bytes. Do not mutate inputs. Raise ValueError for invalid input. This models packed caches, not padding to the batch maximum.

Examples

EXAMPLE 1

Inputlengths=[2,3,1], layers=1, kv_heads=1, head_dim=2, bytes_per_element=2, budget_bytes=40

Outputsizes=[16,24,8], admitted=[0,1], used=40, remaining=0

Exact budget fits.
EXAMPLE 2

Inputlengths=[6,1], same dimensions, budget_bytes=40

Outputadmitted=[]

Arrival-order policy stops at the first oversized request.
solution.pyPython 3.12