Token-Budgeted Chunked Prefill

Problem

Schedule one prefill iteration over requests in their existing order. Visit each request at most once. For every unfinished request, assign min(its remaining tokens, max_chunk, remaining iteration budget), skipping zero assignments.

Return {"chunks": [request_index,token_count] pairs, "remaining": updated per-request counts, "unused": unspent token budget}. Do not revisit the first request after a later one finishes; a per-request chunk cap may leave budget unused.

def chunked_prefill(remaining_tokens, token_budget, max_chunk):

Constraints

• remaining_tokens contains 0..32 integers in [0,128].
• 0 ≤ token_budget ≤128; 1 ≤ max_chunk ≤32.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    5,
    3,
    2
  ],
  6,
  4
]
Expected result: {
  "chunks": [
    [
      0,
      4
    ],
    [
      1,
      2
    ]
  ],
  "remaining": [
    1,
    1,
    2
  ],
  "unused": 0
}

Arguments: [
  [
    10
  ],
  8,
  3
]
Expected result: {
  "chunks": [
    [
      0,
      3
    ]
  ],
  "remaining": [
    7
  ],
  "unused": 5
}